Using LKJCorr together with MvNormal

Questions version agnostic
30

Can you post a MVE that reproduces the bug? It samples fine if I am just trying to estimate the correlations, like this:

import pymc as pm
import pytensor.tensor as pt
import numpy as np
import arviz as az

# Generate Data
n = 5
eta = 1

corr_values = pm.LKJCorr.dist(n=n, eta=eta)
corr = pt.zeros((n, n))
corr = pt.set_subtensor(corr[np.triu_indices(n, k=1)], corr_values)
corr = corr + corr.T + pt.identity_like(corr)

chol_corr = pt.linalg.cholesky(corr)

chol = pm.draw(chol_corr, 1)
y = pm.draw(pm.MvNormal.dist(mu=0, chol=chol), 100)

# PyMC Model
with pm.Model() as m:
    
    corr_values = pm.LKJCorr('corr_values', n=n, eta=eta)
    corr = pt.zeros((n, n))
    corr = pt.set_subtensor(corr[np.triu_indices(n, k=1)], corr_values)
    corr = corr + corr.T + pt.identity_like(corr)

    chol_corr = pt.linalg.cholesky(corr)
    y_hat = pm.MvNormal('y_hat', mu=0, chol=chol_corr, observed=y)
    idata = pm.sample()

Without seeing a model, my hypothesis is that the sampler is proposing illegal values that should just be discarded because they have -np.inf logp, but are causing a computation problem for you before that can happen (see here for a similar situation in a totally unrelated context). If I’m right, the fix is just to pass on_error='nan' to pt.linalg.cholesky. That will prevent it from erroring out if a proposal is not PSD, and the proposal will be rejected anyway so it’s no harm no foul.