Hidden Markov model in pyro and numpyro

I took this great unsupervised example written in pyro and translated it into numypro (jax).

Also check this article how to do this in pymc