Thanks for the help!
Unfortunately my model does not compile with any of those (nutpie, blackjax, numpyro) and I haven’t found a way to debug the problem. It likely has to do with some complicated shape issue / the scan operation. It runs fine with pm.sample and parameter recovery even worked well with a (much) smaller dataset than the one I have (though variational inference did not - the posterior for the last component of a dirichlet was systematically underestimated).