Why does MAP vastly outperform sample in bayesian clustering?

I ran the model on a 2021 macbook pro m1, it took about 10 minutes in total (4 minutes ADVI, 6 minutes sampling). JAX took about 4 minutes in total but I got worse traces.

The fact that MvNormal specifically is giving you grief makes me worry about your BLAS installation. If you’re on an M1/M2 mac make sure you’re using accelerate, otherwise make sure you’re using MKL. If you’re on Linux, slow sampling rates on MvNormal models have been reported on Ubuntu: see here, and the associated issue here.