oops, no, I get the same error messageModuleNotFoundError: No module named 'jax.experimental'; 'jax' is not a package.
This means I should go straight to Jax support, right?
EDIT: found another jax.py in the PATH that was interfering, sorry for the noise!
I’ll mark it as solved.
Thank you as always Ricardo!
I will try the new method you mentioned of using nuts_sampler="numpyro"