# Max\_steps with sampling\_jax.sample\_numpyro\_nuts

**URL:** <https://discourse.pymc.io/t/max-steps-with-sampling-jax-sample-numpyro-nuts/10119>\
**Category:** v5\
**Tags:** sampling\
**Created:** [August 10, 2022, 12:22pm UTC](https://discourse.pymc.io/t/max-steps-with-sampling-jax-sample-numpyro-nuts/10119 "2022-08-10T12:22:04Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![Hans](https://avatars.discourse-cdn.com/v4/letter/h/e19b73/32.png) [@Hans](https://discourse.pymc.io/u/Hans)\
**Post date:** [August 10, 2022, 12:22pm UTC](https://discourse.pymc.io/t/max-steps-with-sampling-jax-sample-numpyro-nuts/10119/1 "2022-08-10T12:22:04Z")

</div>

I have a hierarchical regression model with two levels, countries and local areas. The country level has 84 different ids and there are 156.000 local areas. The total number of rows in the full dataset is about 2.500.000. When I fit the model to a subset of the data, all rows that has a local area id in a random sample of 12.000 of the 156.000 local areas, NUTS converges pretty quickly, but as I increase the number of local areas included in the subset, NUTS seems to get stuck with too small step sizes (around 10^-6). My interpretation is that when the number of parameters gets significantly larger than the default max number of steps (1024) NUTS gets in trouble finding a configuration of parameter values that are better than the current state, so it tries to decrease the step size. It seem NUTS does work as long as the number of parameters is below 10 \* max number of steps. Does my interpretation of what happens sound reasonable?

If it is, I think the solution should be to increase the maximum number of steps per iteration. For the full set with 156.000 parameters, the maximum number of steps per iteration should perhaps be something like 15.000 instead of the default 1024. I have tried both setting `max_steps=2048` and `max_tree_size=11`, but it seems these are ignored since the output I see never prints a step size higher than 1023.

```auto
15/500 [02:09<1:33:53, 11.62s/it, 1023 steps of size 6.63e-06. acc. prob=0.59]

```

Here is the function call I use for the sampling

```auto
my_fit_pymc = pm.sampling_jax.sample_numpyro_nuts(random_seed=1234, tune=250, draws=250, target_accept=0.90, chains = 2, chain_method='sequential', idata_kwargs=dict(log_likelihood=False, max_tree_size=11, max_steps=2048))

```

---

<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:** [August 10, 2022, 1:55pm UTC](https://discourse.pymc.io/t/max-steps-with-sampling-jax-sample-numpyro-nuts/10119/2 "2022-08-10T13:55:49Z")

</div>

I have no idea about the theoretical questions.

Those parameters should not be part of the idata\_kwargs, but passed directly to the sample function. However this type of arguments were not being forwarded until very recently so if it’s not working yet in the latest release you’ll have to wait for the next one.

---

<div class="post-metadata">

**Author:** ![Hans](https://avatars.discourse-cdn.com/v4/letter/h/e19b73/32.png) [@Hans](https://discourse.pymc.io/u/Hans)\
**Post date:** [August 10, 2022, 3:32pm UTC](https://discourse.pymc.io/t/max-steps-with-sampling-jax-sample-numpyro-nuts/10119/3 "2022-08-10T15:32:08Z")

</div>

Thanks for your reply! I guess you refer to [Pass user-provided NUTS kwargs to Numpyro by jhrcook · Pull Request #6021 · pymc-devs/pymc · GitHub](https://github.com/pymc-devs/pymc/pull/6021)

I looked over it and as I understood it, it only affects ’adapt\_step\_size’, ’adapt\_mass\_matrix’ and ’dense\_mass’, but I’ll try it out.

---

<div class="post-metadata">

**Author:** ![Hans](https://avatars.discourse-cdn.com/v4/letter/h/e19b73/32.png) [@Hans](https://discourse.pymc.io/u/Hans)\
**Post date:** [August 10, 2022, 4:55pm UTC](https://discourse.pymc.io/t/max-steps-with-sampling-jax-sample-numpyro-nuts/10119/4 "2022-08-10T16:55:35Z")

</div>

```auto

my_fit_pymc = pm.sampling_jax.sample_numpyro_nuts(random_seed=1234, tune=250, draws=250, target_accept=0.90, chains = 2, chain_method='sequential', idata_kwargs=dict(log_likelihood=False), nuts_kwargs=dict(max_tree_depth=12))

```

Allowed me to get more steps per iteration.

---

<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:** [August 11, 2022, 4:05am UTC](https://discourse.pymc.io/t/max-steps-with-sampling-jax-sample-numpyro-nuts/10119/5 "2022-08-11T04:05:28Z")

</div>

Should work for all parameters, but we only tested those explicitly.

---

<div class="post-metadata">

**Author:** ![junpenglao](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/junpenglao/32/8_2.png) [@junpenglao](https://discourse.pymc.io/u/junpenglao)\
**Post date:** [August 11, 2022, 5:07am UTC](https://discourse.pymc.io/t/max-steps-with-sampling-jax-sample-numpyro-nuts/10119/6 "2022-08-11T05:07:50Z")

</div>

> [@Hans](#):
>
> It seem NUTS does work as long as the number of parameters is below 10 \* max number of steps. Does my interpretation of what happens sound reasonable?

It is not guarantee. You situation is most likely some parameters is/are stuck at difficult geometry, you should look at modifying the prior and inspect the posterior before increasing the number of steps.