Last time I looked at it, Python had an awkward situation with with probabilistic programming packages. There's PyMC3 which is fairly robust and feature-complete, but mostly for static models. Simple things are easy, moderate things are (to me) hard. It also had a bit of an uncertain future; PyMC3 was reliant on Theano, then PyMC4 was conceived, based on TensorFlow, then abandoned in favour of supporting Theano, etc etc.

There's Edward and TF Probability, but somehow I never ever got on with it.

Then there's Pyro, backed by Uber, which again seems fairly feature-complete. Simple things were moderately hard, but hard things were doable. It is also based on PyTorch, which is certainly much easier (IMO) to use than TensorFlow, but still marginally harder than numpy. Also last time I looked, MCMC sampling was a bit missing.

NumPyro perhaps fits into the gap in the happy middle? Most of the power of Pyro, using more familiar numpy-esque (JAX) backend, with an emphasis on simplicity. I should give it a try.

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/