`Interpolated` not supported in JAX/Numba backends

[FBruzzesi]

Describe the issue:

I am raising the issue to create awareness, as I wasn't able to find anything about this error.

While trying to replicate the updating priors notebook, I noticed that if we switch sampler to numpyro, then the following error is raised

NotImplementedError: No JAX conversion for the given Op: SplineWrapper{spline=}

Reproduceable code example:

for _ in range(10):
    ...
    with Model():
        ...
        trace = sample(1000, nuts_sampler = "numpyro", nuts_sampler_kwargs = {"chain_method": "parallel"})
        traces.append(trace)

Error message:

NotImplementedError: No JAX conversion for the given Op: SplineWrapper{spline=}

PyMC version information:

pymc version 5.6.1
numpyro version 0.12.1
jax version 0.4.14

[welcome]

🎉 Welcome to PyMC! 🎉 We're really excited to have your input into the project! 💖

If you haven't done so already, please make sure you check out our Contributing Guidelines and Code of Conduct.

[ricardoV94]

Interpolated uses a custom Op (SplineWrapper) that doesn't yet have a JAX (or Numba) implementations.

[leotercas]

Hello FBruzzesi, did you find any alternative solution for this problem?

[FBruzzesi]

Hey @leotercas , not really no 😭