# Open all NUTS kwargs for sampling with Numpyro

**URL:** <https://discourse.pymc.io/t/open-all-nuts-kwargs-for-sampling-with-numpyro/10040>\
**Category:** v5\
**Tags:** development, jax\
**Created:** [July 29, 2022, 12:22pm UTC](https://discourse.pymc.io/t/open-all-nuts-kwargs-for-sampling-with-numpyro/10040 "2022-07-29T12:22:11Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![jhrcook](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/jhrcook/32/3255_2.png) [@jhrcook](https://discourse.pymc.io/u/jhrcook)\
**Post date:** [July 29, 2022, 12:22pm UTC](https://discourse.pymc.io/t/open-all-nuts-kwargs-for-sampling-with-numpyro/10040/1 "2022-07-29T12:22:11Z")

</div>

In the current implementation `pymc.sampling_jax.sample_numpyro_nuts()`, only some of the arguments for the Numpyro NUTS sampler are available to the user because a few are preset in the function:

```python
if nuts_kwargs is None:
    nuts_kwargs = {}
nuts_kernel = NUTS(
    potential_fn=logp_fn,
    target_accept_prob=target_accept,
    adapt_step_size=True,
    adapt_mass_matrix=True,
    dense_mass=False,
    **nuts_kwargs,
)

```

I think it would be useful to make all of these arguments available through the PyMC interface, but before proposing a change I was curious if there was that there are reason they aren’t? If not, would this be a PR I could make?

---

<div class="post-metadata">

**Author:** ![cluhmann](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/cluhmann/32/3083_2.png) [@cluhmann](https://discourse.pymc.io/u/cluhmann)\
**Post date:** [July 29, 2022, 12:33pm UTC](https://discourse.pymc.io/t/open-all-nuts-kwargs-for-sampling-with-numpyro/10040/2 "2022-07-29T12:33:20Z")

</div>

Those are only being “preset” as defaults. They should all be available to be set explicitly by the user, no?

---

<div class="post-metadata">

**Author:** ![jhrcook](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/jhrcook/32/3255_2.png) [@jhrcook](https://discourse.pymc.io/u/jhrcook)\
**Post date:** [July 29, 2022, 12:42pm UTC](https://discourse.pymc.io/t/open-all-nuts-kwargs-for-sampling-with-numpyro/10040/3 "2022-07-29T12:42:27Z")

</div>

Unfortunately, no because if you pass in a argument in `nuts_kwargs` that is already preset, you get an error that looks something like this (example from trying to pass an argument to `dense_mass`):

```python
 numpyro.infer.hmc.NUTS() got multiple values for keyword argument 'dense_mass'

```

---

<div class="post-metadata">

**Author:** ![cluhmann](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/cluhmann/32/3083_2.png) [@cluhmann](https://discourse.pymc.io/u/cluhmann)\
**Post date:** [July 29, 2022, 12:50pm UTC](https://discourse.pymc.io/t/open-all-nuts-kwargs-for-sampling-with-numpyro/10040/4 "2022-07-29T12:50:09Z")

</div>

Ah, I see. So arguments like `adapt_step_size` , `adapt_mass_matrix`, and `dense_mass` aren’t exposed for pass-through in `sample_numpyro_nuts()`. There may be reasons (that I am not aware of), but if you submit an issue, you can figure it out one way or another.

---

<div class="post-metadata">

**Author:** ![jhrcook](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/jhrcook/32/3255_2.png) [@jhrcook](https://discourse.pymc.io/u/jhrcook)\
**Post date:** [July 29, 2022, 12:52pm UTC](https://discourse.pymc.io/t/open-all-nuts-kwargs-for-sampling-with-numpyro/10040/5 "2022-07-29T12:52:12Z")

</div>

Great, will do! Thank you

Link to issue: [Open all NUTS kwargs to user for sampling with Numpyro · Issue #6020 · pymc-devs/pymc · GitHub](https://github.com/pymc-devs/pymc/issues/6020)

---

<div class="post-metadata">

**Author:** ![jhrcook](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/jhrcook/32/3255_2.png) [@jhrcook](https://discourse.pymc.io/u/jhrcook)\
**Post date:** [August 1, 2022, 10:40am UTC](https://discourse.pymc.io/t/open-all-nuts-kwargs-for-sampling-with-numpyro/10040/6 "2022-08-01T10:40:43Z")

</div>

Pull request: [Pass user-provided NUTS kwargs to Numpyro by jhrcook · Pull Request #6021 · pymc-devs/pymc · GitHub](https://github.com/pymc-devs/pymc/pull/6021)
