`SA` kernel fails for chain method `vectorized`
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 by tracing SA through NumPyro's MCMC chain handling for chain_method="vectorized", then compare the failing minimal example with the sequential and parallel cases. Run the reproduction with two scalar chains and confirm the vectorized run completes without the reported vmap rank error.
Written by the indexing model from the issue text.
Description
Bug Description
Sampling with the SA kernel with chain_method="vectorized" leads to a crash in most cases.
Steps to Reproduce
Minimal example:
import jax
import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, SA
def model():
numpyro.sample("x", dist.Normal(0, 1))
mcmc = MCMC(
SA(model),
num_warmup=100,
num_samples=100,
num_chains=2,
chain_method="vectorized",
)
mcmc.run(jax.random.key(0))
This raises
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 ())
This problem doesn't occur for chain method sequential or parallel or if the dimension of x coincides with num_chains.
Versions
numpyro 1e42fd58e5d6d60f8a767915c524edb787f2e81b
jax 0.10.2
Expected Behavior
The sampler to run without crashing.
- 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` with vectorized chains crashes when re-ran with changed parameters (similar to #2338)Possibly taken @tandede claimed this 3 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 10 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 20 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 16 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
-
Difficulty 2/5 1-3 hours Newbie friendliness 68/100
-
Difficulty 2/5 1-3 hours Newbie friendliness 70/100
Maintainers usually reply within 2 days
-
Difficulty 2/5 1-3 hours Newbie friendliness 62/100
Maintainers usually reply within 2 days
-
bug
Difficulty 2/5 1-3 hours Newbie friendliness 82/100
developmentseed/titiler-cmr#278 ·
Maintainers usually reply within 1 day
-
Difficulty 2/5 1-3 hours Newbie friendliness 82/100
Maintainers usually reply within 2 days