These examples are super nice! I’m also doing something to have more explicit examples! See NumPyro time series examples