Will take a look. Interesting story is it works for multivariate case based on a simulation that I did, which makes me wonder how it works in the first place. It was a simple set up where data is from multivariate normal with zero mean and cov matrix 3x3 (manually set as long as it’s pd). I masked each column 30% nan randomly and turned out it not only recovered the cov matrix pretty well (close to true), but delivered better imputation than using col means. The likelihood I used is:
out = pm.MvNormal(‘out’, mu=0, chol=chol, observed=masked_z)
where chol has LKJ prior and I was using ADVI to do inference. Any idea?