@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 doesn’t work: import pymc as pm import numpy as np with pm.Model(): a = pm.Normal("a") idata = pm.sample(nuts_sampler…

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 CC @fonnesbeck

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

Questions v5

pymcmercury June 9, 2023, 5:59pm 4

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

Topic		Replies	Views
How can I set the maximum tree depth for the NUTS method from the blackjax library? v5	0	321	June 9, 2023
NUTS slows dramatically after exactly 200 samples Questions	2	738	August 25, 2019
Open all NUTS kwargs for sampling with Numpyro v5 development , jax	5	823	August 1, 2022
Max_steps with sampling_jax.sample_numpyro_nuts v5 sampling	5	584	August 11, 2022
What does the "maximum tree depth" warning mean? Questions	6	7203	December 4, 2019

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

Related topics