I've been using Numpyro for a while (2 years), and it's very fast, robust and well designed.

I recommend the Statistical rethinking translation [0].

I've had some trouble with the plates sometimes, especially when using discrete latent variables.

With PyMC these are the best for bayesian modelling.

[0] https://fehiepsi.github.io/rethinking-numpyro/