How do I manipulate shape when using sample_numpyro_nuts

[cluhmann]

What version of v4 are you running? I installed (at some point) straight from the repo and your code works fine for me.