Difference between jax.numpy.linalg.cholesky, jax.scipy.linalg.cholesky & jax.lax.linalg.cholesky

[adam-hartshorne]

What is the difference between these 3 functions?

[hawkinsp]

• jax.lax.linalg.cholesky: this is JAX's internal cholesky primitive. The other two functions call this internally. We recommend against using lax functions if there's a jax.numpy or jax.scipy equivalent since the latter two APIs are more stable.
• jax.numpy.linalg.cholesky: this API mimics numpy.linalg.cholesky (https://numpy.org/doc/stable/reference/generated/numpy.linalg.cholesky.html) and is implemented using lax.linalg.cholesky (https://github.com/google/jax/blob/fd3a72dd1fc1112b698c9f74ee9720d78280e98d/jax/_src/numpy/linalg.py#L43).
• jax.scipy.linalg.cholesky: this API mimics scipy.linalg.cholesky (https://docs.scipy.org/doc/scipy/reference/generated/scipy.linalg.cholesky.html) and is also implemented using lax.linalg.cholesky (https://github.com/google/jax/blob/fd3a72dd1fc1112b698c9f74ee9720d78280e98d/jax/_src/scipy/linalg.py#L41).

If you're happy with the NumPy API, I'd pick that one.