`SA` with vectorized chains crashes when re-ran with changed parameters (similar to #2338)
Maintainers usually reply within 1 day
Assessment
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Newbie friendliness
- 65/100
- Issue type
- Bug
- Clarity
- Mostly clear
- Activity status
- Active
- Tech stack
- python
- Domain
- machine-learning
Research direction
Start in NumPyro's MCMC and SA implementation, then compare how vectorized chains handle repeated runs with changed model parameters against the related HMC issue #2338 and discussion in #2339. Reproduce the supplied example and trace the vmap argument shape on the second mcmc.run; done when both runs complete with the changed parameter and vectorized chains.
Written by the indexing model from the issue text.
Description
Description
When using the SA 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 ())
SA shows the same behavior as HMC in #2338, likely also with a similar cause (see discussion of #2339).
Minimal Example
from jax import random
import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, SA
def model_dynamic_param(mean):
numpyro.sample("x", dist.Normal(mean, 1).expand([2]))
kernel = SA(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)
- Dominant language
- Python
- Stars
- 2.8k
- Forks
- 338
- Avg merge
- 2d 8h
- Merged PRs (30d)
- 67
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
-
`SA` kernel fails for chain method `vectorized`Possibly taken @tandede claimed this 4 days ago. Openbug
Difficulty 3/5 1-2 days Newbie friendliness 65/100
Maintainers usually reply within 1 day
-
contrib.hsgp: hsgp_periodic_non_centered scales beta by the variance coefficient instead of its square rootPossibly taken @Vamsi-2581 claimed this 12 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
-
HMC with vectorized chains crashes when re-ran with changed parametersPossibly taken @ingolfuer claimed this 21 days ago. Openbug
Difficulty 2/5 1-3 hours Newbie friendliness 50/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 18 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
All issues in pyro-ppl/numpyro
Similar issues
-
bug
Difficulty 2/5 1-3 hours Newbie friendliness 74/100
BCDA-APS/polar-bits#93 ·
-
infer_audio_type only recognises one MPEG Layer III frame header, so most MP3 files are rejectedPossibly taken @Steve-20-ai claimed this today. Open
Difficulty 2/5 1-3 hours Newbie friendliness 78/100
Maintainers usually reply within 1 day
-
[BUG] - ssh_config `Match tagged` statement always appliesPossibly taken @AtharvRG claimed this today. OpenBug
Difficulty 2/5 1-3 hours Newbie friendliness 68/100
paramiko/paramiko#2715 · 1 comment ·
Maintainers usually reply within 9 days
-
auto-rejected intake mcp-intake needs-salvage priority:high type:feature
Difficulty 2/5 1-3 hours Newbie friendliness 65/100
Ikalus1988/MisakaNet#4221 · 2 comments ·
Maintainers usually reply within 1 day
-
Difficulty 2/5 1-3 hours Newbie friendliness 82/100
Maintainers usually reply within 1 day