If you mean this example (a random walk), I have a notebook here with more examples that work case_studies/Time Series with Numpyro - Random Walk Intro.ipynb at main · kylejcaron/case_studies · GitHub
If you meant the HMM I linked, I was just referring to the pattern they use of poutine.mask (numpyros equivalent is handlers.mask) - cant speak for the rest of their code so sorry if that confused you.
Here are better examples I just found from searching for handlers.mask and poutine.mask Variable length sequences and masking Example: Hidden Markov Models — Pyro Tutorials 1.9.1 documentation