# How can I set the maximum tree depth for the NUTS method from the numpyro library?

**URL:** <https://discourse.pymc.io/t/how-can-i-set-the-maximum-tree-depth-for-the-nuts-method-from-the-numpyro-library/12258>\
**Category:** v5\
**Created:** [June 5, 2023, 2:39pm UTC](https://discourse.pymc.io/t/how-can-i-set-the-maximum-tree-depth-for-the-nuts-method-from-the-numpyro-library/12258 "2023-06-05T14:39:46Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![pymcmercury](https://avatars.discourse-cdn.com/v4/letter/p/7cd45c/32.png) [@pymcmercury](https://discourse.pymc.io/u/pymcmercury)\
**Post date:** [June 5, 2023, 2:39pm UTC](https://discourse.pymc.io/t/how-can-i-set-the-maximum-tree-depth-for-the-nuts-method-from-the-numpyro-library/12258/1 "2023-06-05T14:39:46Z")

</div>

@ricardoV94 @junpenglao How can I set the maximum tree depth for the NUTS method from the `numpyro` library?  
The way described in the test file [test\_mcmc\_external.py](https://github.com/pymc-devs/pymc/blob/da68d11f8bdddf4d3730e489fc4046ee101f850b/tests/sampling/test_mcmc_external.py) doesn’t work:

```auto
import pymc as pm
import numpy as np

with pm.Model():
        a = pm.Normal("a")
        idata = pm.sample(nuts_sampler = "numpyro",
                          target_accept = 0.99,
                          nuts = {"max_treedepth": 1},
                          random_seed = 1410)

print(np.max(idata.sample_stats.tree_depth))
# <xarray.DataArray 'tree_depth' ()>
# array(4)

```

and specifying something via the `nuts_kwargs` argument throws `ValueError: Unused step method arguments: {'nuts_kwargs'}`.

---

<div class="post-metadata">

**Author:** ![ricardoV94](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/ricardov94/32/5775_2.png) [@ricardoV94](https://discourse.pymc.io/u/ricardoV94)\
**Post date:** [June 7, 2023, 7:45am UTC](https://discourse.pymc.io/t/how-can-i-set-the-maximum-tree-depth-for-the-nuts-method-from-the-numpyro-library/12258/2 "2023-06-07T07:45:52Z")

</div>

Yeah, I don’t see any code path to pass any arguments other than `target_accept`, I guess you could go directly to the `numpyro_sampler` function and pass it as `nuts_kwargs`? [pymc/jax.py at 261862d778910a09c5b61edcc66958519a86815e · pymc-devs/pymc · GitHub](https://github.com/pymc-devs/pymc/blob/261862d778910a09c5b61edcc66958519a86815e/pymc/sampling/jax.py#L517)

CC @fonnesbeck

---

<div class="post-metadata">

**Author:** ![ricardoV94](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/ricardov94/32/5775_2.png) [@ricardoV94](https://discourse.pymc.io/u/ricardoV94)\
**Post date:** [June 7, 2023, 7:59am UTC](https://discourse.pymc.io/t/how-can-i-set-the-maximum-tree-depth-for-the-nuts-method-from-the-numpyro-library/12258/3 "2023-06-07T07:59:53Z")

</div>

Opened an issue here: [Numpyro sampler `nuts_kwargs` can't be passed from `sample` · Issue #6757 · pymc-devs/pymc · GitHub](https://github.com/pymc-devs/pymc/issues/6757)

---

<div class="post-metadata">

**Author:** ![pymcmercury](https://avatars.discourse-cdn.com/v4/letter/p/7cd45c/32.png) [@pymcmercury](https://discourse.pymc.io/u/pymcmercury)\
**Post date:** [June 9, 2023, 5:59pm UTC](https://discourse.pymc.io/t/how-can-i-set-the-maximum-tree-depth-for-the-nuts-method-from-the-numpyro-library/12258/4 "2023-06-09T17:59:52Z")

</div>

Yes, passing it directly to `pymc.sampling.jax.sample_numpyro_nuts` via the `nuts_kwargs` argument works. Thanks!