BUG: pm.sample can't disable init-value jitter for the external JAX samplers (numpyro/blackjax)

Start at pm.sample's external-sampler path and _sample_external_nuts, then inspect how sample_jax_nuts receives its jitter argument. Reproduce the issue with the provided numpyro example and verify that jitter control reaches sample_jax_nuts rather than the NUTS kernel, without the reported TypeError.

Describe the issue:

pymc.sampling.jax.sample_jax_nuts has a jitter: bool = True argument, but one can't reach it through pm.sample(nuts_sampler="numpyro"|"blackjax"). _sample_external_nuts calls sample_jax_nuts(...) without forwarding jitter, so it's always True, and passing it via nuts= / nuts_sampler_kwargs lands in the NUTS kernel instead:

import pymc as pm
with pm.Model() as m:
    pm.Normal("x", 0, 1)
    pm.sample(nuts_sampler="numpyro", nuts={"jitter": False})
# TypeError: NUTS.__init__() got an unexpected keyword argument 'jitter'
The only way to set it is calling pymc.sampling.jax.sample_jax_nuts(jitter=False) directly, i.e. bypassing pm.sample.

Could pm.sample forward a jitter control to sample_jax_nuts for the external samplers (or have init="adapt_diag" imply jitter=False)? The argument already exists — it just isn't plumbed through.

Note: the in-process pymc sampler can turn init jitter off via init="adapt_diag" (I realize the default is "adapt_diag+jitter", but the JAX path can't be matched — init doesn't govern it (pm.sample even warns init='adapt_diag' is ignored by nuts_sampler='numpyro'). When you supply your own carefully-computed initial values and don't want the extra uniform[-1, 1] jitter on top, that's being made hard currently.

PyMC version information: PyMC 6.1.0 · Python 3.12 (also affects blackjax).

P.S.: I realize this is not strictly speaking a bug but considering it a new feature seems wrong to. Seems to me more like a regression that is being asked to be undone.

Reproduceable code example:

import pymc as pm
with pm.Model() as m:
    pm.Normal("x", 0, 1)
    pm.sample(nuts_sampler="numpyro", nuts={"jitter": False})
The only way to set it is calling pymc.sampling.jax.sample_jax_nuts(jitter=False) directly, i.e. bypassing pm.sample.

Error message:

# TypeError: NUTS.__init__() got an unexpected keyword argument 'jitter'

PyMC version information:

PyMC version information: PyMC 6.1.0 · Python 3.12 (also affects blackjax).

Context for the issue:

I am working on a toolbox for which controlling starting points (and the scale of the jitter in particular) is relevant in some contexts.