pymc-devs/pymc

Use same API for defining internal and external Nuts kwargs

Open

#6,757 opened on Jun 7, 2023

View on GitHub
 (2 comments) (0 reactions) (0 assignees)Python (1,902 forks)batch import
bughelp wantedjax

Repository metrics

Stars
 (7,926 stars)
PR merge metrics
 (Avg merge 11d) (7 merged PRs in 30d)

Description

Description

User on discourse reported:

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 = "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'}.

I don't know if nuts should be converted to nuts_kwargs, but even if a user were to pass nuts_kwargs to sample, those wouldn't make it to the sample_numpyro_nuts function because we drop arbitrary kwargs passed here:

https://github.com/pymc-devs/pymc/blob/261862d778910a09c5b61edcc66958519a86815e/pymc/sampling/mcmc.py#L252

Contributor guide

Research direction
Inspect how kwargs are handled in the sample() call and passed to sample numpyro nuts, specifically around line 252 of mcmc.py. Ensure that the 'nuts' parameter (or 'nuts kwargs') is correctly forwarded and not dropped.
Tech stack
python
Domain
backend
Issue type
Bug
DifficultyEstimated implementation difficulty for a new contributor, from 1 for very small changes to 5 for expert-level work.
2
Estimated timeA rough time range for an experienced contributor to investigate, implement, test, and prepare a pull request.
1-3 hours
Activity statusHow available the issue appears right now: fresh, active, stale, blocked, or waiting on maintainer input.
Active
ClarityHow clearly the issue explains the expected change, acceptance criteria, and next step.
Mostly clear
Prerequisites
PythonPyMC
Newbie friendlinessA 1-100 score estimating how approachable this issue is for first-time contributors.
55