Hello, I have a model with a multivariate Gaussian likelihood whose covariance matrix is not constant, depends on the modelled variables and can change arbitrarily at each step (i.e. not a simple scaling etc.). In high-dimensional cases (D >> 1000) both the construction and inverse/Cholesky decomposition of the covariance matrix become expensive and thus unfeasible at each HMC step.

Is it possible with NumPyro to, for example, only update the covariance matrix every N steps?

I have previously tried attempts where I run the HMC for N steps, update the covariance matrix, run for another N steps and repeat. However, that means that each time I had to completely restart the HMC.

The advantage of the model is that while the covariance matrix depends on the modelled parameters, it does not vary much once the model is converged.

Thank you!