Opened an issue here: Numpyro sampler `nuts_kwargs` can't be passed from `sample` · Issue #6757 · pymc-devs/pymc · GitHub