How do I manipulate shape when using sample_numpyro_nuts

Hi @qipengchen, I’m having the same issue. Did you happen to find a solution to this problem?