Sampling issue with `jax.random.multivariate_normal`

Sampling from a GP is giving unexpected results. For example, consider following code:

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import tinygp

N = 1024

fig, ax = plt.subplots(3,3,figsize=(6.4*1.5, 4.8*1.5))
for seed in range(9):
  key = jax.random.PRNGKey(seed)
  x = 5 * jax.random.normal(key, shape=(N,1))
  kernel = 4 * tinygp.kernels.ExpSquared(scale=0.5)
  y = tinygp.GaussianProcess(kernel, x, diag=1.0).sample(key)
  ax.ravel()[seed].scatter(x, y, s=5);

This code generates the following output:

If I replace tinygp.GaussianProcess(kernel, x, diag=1.0).sample(key) with jax.random.multivariate_normal(key, jnp.zeros(N), kernel(x, x)+jnp.eye(N)), I get the same result. But, if I change the sampling method as jax.random.multivariate_normal(key, jnp.zeros(N), kernel(x, x)+jnp.eye(N), method='svd'), I get drastically different results:

[patel-zeel]

Sure, thank you!

[dfm]

I should say: feel free to open an issue on the JAX repo if you don’t like that, but this isn’t a tinygp question!