Just to rule out something obvious, do you get the same weird results if you just run with vanilla pm.sample, instead of trying to initialize NUTS manually?
Just to rule out something obvious, do you get the same weird results if you just run with vanilla pm.sample, instead of trying to initialize NUTS manually?