Without looking at it, you can try passing check_bounds=False to the Model if you know the sampler will not propose invalid values.
Otherwise there are many things that might be different in the PyMC JAX graph and the NumPyro one. You can try to inspect those