The modeling code looks great! It might be a bit clearer if we put obs... variables into a numpyro.plate context.
obs...
numpyro.plate