multivariate_normal with positive semi-definite matrix and default Cholesky solver returns NaN

[marisgg]

Hi!

Thank you for your work on Jax. Please excuse me if I made any false assumptions or mistakes. I'm very new to Jax. I've came across unexpected NaN's in quite a simple multivariate_normal call. See the following example:

import jax
import jax.numpy as jnp
mu = jnp.zeros((2,))
sigma = jnp.ones((2, 2))
key = jax.random.PRNGKey(0)
jax.random.multivariate_normal(key, mean=mu, cov=sigma)

which gives:

DeviceArray([nan, nan], dtype=float32)

Running with either method='eigh' or method='svd' does produce a valid result without NaN. Cholesky decomposition is the default solver in the function call.

The NaN's appear in the Cholesky function call. It's most probably because the sigma matrix is not positive semi-definite. Shouldn't this be detectable, such that the multivariate_normal method call can throw an error and suggest (or fallback, but that might not be desirable) other solver methods instead of returning NaN?

Otherwise, it could be similar to how numpy.linalg.cholesky throws an error if the input matrix is not positive semi-definite, while jax.numpy.linalg.cholesky returns NaN. Then, if you call multivariate_normal with the default Cholesky solver and a covariance that is not positive semi-infinite an error would be raised in the Cholesky call indicating the error of the combination of covariance and solver method.

I understand that it might be up to the developer to notice these kind of things, but it seems reasonable to raise an error with some explanation rather than returning NaN.

The code example given at the top was ran on jax==0.3.25 with identical jaxlib version on the CPU (Ubuntu 20.04.5 WSL2).

[jakevdp]

Thanks for the report. Perhaps we should default to a method other than cholesky.

Regarding detecting problematic inputs at runtime: this is not necessarily easy in JAX, for the same reason that it's not easy to detect out-of-bound indices at runtime. I won't say more here because the background is pretty well covered in https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#out-of-bounds-indexing

[marisgg]

Thanks for the answer. I hadn't thought of it that way since I am quite new to JAX. That page explains a lot and is a very interesting read altogether. I agree that something other than cholesky might be a better default if it's less strict with respect to its inputs, but I guess that's a design choice.

[jakevdp]

Agreed a different method would be better as a default, but this would be a breaking change because it could change the stream of values produced between JAX releases. I'll add documentation for now, and perhaps we can plan to change it as part of the upcoming overhaul of the random submodule.

[jakevdp]

Doc changes in #13518