HMC with vectorized chains crashes when re-ran with changed parameters
Maintainers usually reply within 1 day
Assessment
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Newbie friendliness
- 50/100
Research direction
The bug is in HMC.init in infer/hmc.py (around line 919), where self._sample_fn is reassigned to a vmapped version on each call. Run the minimal example with chain_method="vectorized" and call mcmc.run twice to reproduce the ValueError. Done means the second run succeeds with a different parameter, without nested vmapping. The reporter has described a fix (a non-vmapped _sample_fn_base) and says a pull request is coming, so check for one before starting.
Written by the indexing model from the issue text.
Description
Description
When using the NUTS sampler with chain_method="vectorized", multiple runs with the same MCMC instance lead to a vmap error that crashes the run:
ValueError: vmap was requested to map its argument along axis 0, which implies that its rank should be at least 1, but is only 0 (its shape is ())
Minimal Example
import jax
from jax import random
import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS
def model_dynamic_param(mean):
numpyro.sample("x", dist.Normal(mean, 1))
kernel = NUTS(model_dynamic_param)
mcmc = MCMC(
kernel, num_warmup=100, num_samples=100, num_chains=2, chain_method="vectorized"
)
mcmc.run(random.key(0), 0)
# different parameter, second run crashes
mcmc.run(random.key(0), 10)
Cause
The cause seems to be in infer/hmc.py, HMC.init overwriting self._sample_fn with a vmapped version of self._sample_fn (line 919):
init_state = vmap(hmc_init_fn)(init_params, rng_key)
assert self._sample_fn is not None
sample_fn = vmap(self._sample_fn, in_axes=(0, None, None))
self._sample_fn = sample_fn
This can happen repeatedly. The resulting function will expect the wrong dimensionality for its inputs, leading to the crash.
Fix
I'm currently writing a fix that introduces _sample_fn_base to the HMC class which stores the non-vmapped version. The vmapped function will always be derived from that, preventing nested vmapping. I'll open a pull request shortly.
- Dominant language
- Python
- Stars
- 2.8k
- Forks
- 334
- Avg merge
- 2d 15h
- Merged PRs (30d)
- 71
Getting set up
- No Dockerfile or Docker Compose file
- No pull request template
- Read the contributing guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
More from pyro-ppl/numpyro
-
contrib.hsgp: hsgp_periodic_non_centered scales beta by the variance coefficient instead of its square rootPossibly taken @Vamsi-2581 claimed this 4 days ago. Openbug
Difficulty 4/5 3-5 days Newbie friendliness 57/100
Maintainers usually reply within 1 day
-
enhancement
Difficulty 5/5 Over a week Newbie friendliness 12/100
Maintainers usually reply within 1 day
-
ZeroInflatedProbs excludes 0 from its support when the base support does not contain 0Possibly taken @GreedyC claimed this 10 days ago. Open
Difficulty 4/5 3-5 days Newbie friendliness 48/100
pyro-ppl/numpyro#2295 · 2 comments ·
Maintainers usually reply within 1 day
-
Port HMM from PyroOpenenhancement
Difficulty 4/5 3-5 days Newbie friendliness 25/100
pyro-ppl/numpyro#2278 · 5 comments ·
Maintainers usually reply within 1 day
-
Easier Use of Distributions Under `vmap`May be free again @DanWaxman claimed this 35 days ago, and no pull request is open. Openenhancement
Difficulty 5/5 Over a week Newbie friendliness 35/100
pyro-ppl/numpyro#2271 · 3 comments · 4 reactions ·
Maintainers usually reply within 1 day
All issues in pyro-ppl/numpyro
Similar issues
-
push_gh-pages fails on master branch on forksPossibly taken A pull request linked to this issue is open or already merged. Opengood first issue infrastructure
Difficulty 2/5 1-3 hours Newbie friendliness 76/100
kuznia-rdzeni/coreblocks#1135 ·
-
good first issue server
Difficulty 2/5 1-3 hours Newbie friendliness 85/100
Maintainers usually reply within 1 day
-
documentation priority: high
Difficulty 2/5 1-3 hours Newbie friendliness 78/100
a-lex-s/mealie-telegram-bot#54 ·
Maintainers usually reply within 1 day
-
type/test
Difficulty 2/5 1-3 hours Newbie friendliness 76/100
SaltedTan/virgil-language-server#218 ·
Maintainers usually reply within 1 day
-
Difficulty 2/5 1-3 hours Newbie friendliness 83/100
btclib-org/btclib-wallet#541 ·
Maintainers usually reply within 1 day