How do I manipulate shape when using sample_numpyro_nuts

Hey @qipengchen, were you able to solve your problem? I have a simillar one.