Check
az.plot_energy
Is a truncated marginal energy distribution a problem? Otherwise looks pretty close to me.
The truncation is also visible with the naive model so I guess not? The naive model is plenty fast.
Also check
az.plot_pair
Well, it’s not completely white noise but cannot complain.
I’d say that there’s more correlation with the naive model.
Also check the number of integration steps
The mean is a bit over 2 and the maximum 4
try switching the nuts sampler
I’m already using blackjax. The default sampler was unusable as even a single sample took ~5 minutes