BUG: pm.sample can't disable init-value jitter for the external JAX samplers (numpyro/blackjax)
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.
Source: pymc-devs/pymc