My apologies, I should’ve read your question more carefully before answering. As you mentioned the problem is with the initial values you assigned. I can run the model with no issues using the default arguments of the sampler:
with pm.Model(coords=coords,) as model:
chol, corr, stds = pm.LKJCholeskyCov(
"chol", n=2, eta=2.0, sd_dist=pm.Exponential.dist(1, shape=2)
)
cov = pm.Deterministic("cov", chol.dot(chol.T), dims=("axis", "axis_bis"))
mu = pm.Normal("mu", 0.0, sigma=1.5, dims="axis")
obs = pm.MvNormal("obs", mu, chol=chol, observed=x, dims=("obs_id", "axis"))
with model:
idata = pm.sample(1000)
az.summary(idata)
Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [chol, mu]
|████████████| 100.00% [8000/8000 00:26<00:00 Sampling 4 chains, 0 divergences]Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 46 seconds.
'''
mean sd hdi_3% ... ess_bulk ess_tail r_hat
mu[y] 0.995 0.007 0.983 ... 3905.0 3155.0 1.0
mu[z] -1.977 0.015 -2.007 ... 3709.0 3044.0 1.0
chol[0] 0.708 0.005 0.699 ... 4736.0 3434.0 1.0
chol[1] -0.606 0.014 -0.632 ... 4359.0 3084.0 1.0
chol[2] 1.362 0.010 1.345 ... 5933.0 3117.0 1.0
chol_corr[0, 0] 1.000 0.000 1.000 ... 4000.0 4000.0 NaN
chol_corr[0, 1] -0.406 0.008 -0.422 ... 4562.0 3382.0 1.0
chol_corr[1, 0] -0.406 0.008 -0.422 ... 4562.0 3382.0 1.0
chol_corr[1, 1] 1.000 0.000 1.000 ... 3124.0 3434.0 1.0
chol_stds[0] 0.708 0.005 0.699 ... 4736.0 3434.0 1.0
chol_stds[1] 1.491 0.011 1.472 ... 5046.0 3070.0 1.0
cov[y, y] 0.501 0.007 0.488 ... 4736.0 3434.0 1.0
cov[y, z] -0.429 0.012 -0.449 ... 3899.0 3118.0 1.0
cov[z, y] -0.429 0.012 -0.449 ... 3899.0 3118.0 1.0
cov[z, z] 2.223 0.031 2.167 ... 5046.0 3070.0 1.0
Results seem sensible, maybe you don’t really need to assign initial values. Is there a specific reason for using them?