Provide jax.scipy.stats.multivariate_normal with a cholesky-factored covariance?

[ahwillia]

I have code where I frequently call jax.scipy.stats.multivariate_normal.logpdf(x, mean, cov) with different values of x but with constant mean and cov. Obviously, it would be much more efficient if I could computer the cholesky factorization of cov once up front rather than have logpdf calculate it repeatedly under the hood. Maybe I'm missing it but it doesn't seem possible to cache the factored matrix? Is this outside the scope of what jax is trying to provide?

[sharadmv]

Insofar as scipy.stats doesn't provide this functionality, I think it is out of scope for JAX for now.

However, consider using TFP on JAX's MVNChol implementation.

import tensorflow_probability.substrates.jax as tfp
import jax.numpy as jnp
tfd = tfp.distributions
mean = jnp.zeros(5)
cov_chol = jnp.linalg.cholesky(jnp.eye(5))

dist = tfd.MultivariateNormalTriL(loc=mean, scale_tril=cov_chol)
dist.sample(seed=...)
dist.log_prob(...)