ZeroInflatedProbs excludes 0 from its support when the base support does not contain 0
Maintainers usually reply within 1 day
Assessment
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Newbie friendliness
- 48/100
- Issue type
- Bug
- Clarity
- Mostly clear
- Activity status
- Active
- Tech stack
- python
- Domain
- machine-learning
Research direction
Start by reading the ZeroInflatedProbs and ZeroInflatedLogits implementations, especially support, enumerate_support, and log_prob, then run the reproduction against a DiscreteUniform base with low greater than 0. Trace how validation and enumeration use those entry points. Done means support and enumeration include 0, log_prob(0) is finite according to gate, and enumerated probabilities sum to 1 without breaking inference.
Written by the indexing model from the issue text.
Description
Bug Description
ZeroInflatedProbs (and ZeroInflatedLogits) forward support and enumerate_support to the base distribution, so when the base support does not contain 0 the zero-inflated distribution has a support without 0, although it samples 0 with probability gate. log_prob is decorated with @validate_sample, so log_prob(0) is -inf: the value it samples most often has probability 0.
This came up in the review of #2290 (discrete samplers and supports that do not start at 0). It affects any base whose support excludes 0, e.g. DiscreteUniform(low, high) with low > 0, or a custom zero-truncated count distribution used to build a hurdle-style model. Poisson, Binomial and the other built-in count bases include 0, so they are not affected.
Steps to Reproduce
On master (ca122297), with JAX 0.11.2:
import jax, jax.numpy as jnp
import numpyro, numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS, DiscreteHMCGibbs
d = dist.ZeroInflatedProbs(dist.DiscreteUniform(3, 7), gate=0.4)
s = d.sample(jax.random.key(0), (10000,))
print((s == 0).mean()) # 0.40
print(d.support.check(0)) # False
print(d.enumerate_support()) # [3 4 5 6 7]
print(jnp.exp(d.log_prob(d.enumerate_support())).sum()) # 0.6
print(jnp.exp(d.log_prob(0))) # 0.0, with an out-of-support warning
def model():
numpyro.sample("x", dist.ZeroInflatedProbs(dist.DiscreteUniform(3, 7), gate=0.4))
numpyro.sample("y", dist.Normal())
MCMC(DiscreteHMCGibbs(NUTS(model)), num_warmup=300, num_samples=3000).run(jax.random.key(1))
# RuntimeError: Cannot find valid initial parameters. Please check your model again.
Enumeration-based inference (TraceEnum_ELBO, infer_discrete) sums over enumerate_support, which lacks 0, and the discrete Gibbs samplers can never accept 0 since log_prob(0) is -inf.
Expected Behavior
The support and enumerate_support include 0, log_prob(0) equals log(gate + (1 - gate) * base.prob(0)) (just log(gate) when 0 is outside the base support), and the probabilities over enumerate_support sum to 1.
The existing constraints can't express a support like {0} ∪ {3, ..., 7} directly (integer_interval(0, 7) would also admit 1 and 2, where DiscreteUniform.log_prob doesn't check the support), so the fix probably needs either a union-style constraint or a support check inside log_prob. Happy to help with a PR once there's a preferred direction.
- 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 3 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 12 days ago. Openbug
Difficulty 2/5 1-3 hours Newbie friendliness 50/100
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 34 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
-
Difficulty 2/5 1-3 hours Newbie friendliness 72/100
CLIMADA-project/climada_python#1337 ·
Maintainers usually reply within 3 days
-
tail without reverseOpen
Difficulty 2/5 1-3 hours Newbie friendliness 68/100
8go/matrix-commander#233 ·
-
documentation
Difficulty 2/5 1-3 hours Newbie friendliness 66/100
-
documentation enhancement priority:high size:M
Difficulty 2/5 1-3 hours Newbie friendliness 82/100
Maintainers usually reply within 1 day
-
Validate non-divisible effective batch sizes in the Unsloth LoRA backendPossibly taken @debesh-biswas claimed this today. Open
Difficulty 2/5 1-3 hours Newbie friendliness 70/100