Do you know how to do this if each of your random walks have a different number of data points?
I think you could probably do something like this with numpyro.handlers.mask and pad with nulls Time varying dimension measurement in gaussian hmm
if you search handlers.mask in this forums search there seem to be some similar examples that might help