Implement icdf for Poisson, Binomial, and NegativeBinomial by jacobpstein · Pull Request #8339 · pymc-devs/pymc

jacobpstein · 2026-06-22T23:13:06Z

Description

Adds icdf (quantile / inverse-CDF) methods for four discrete distributions:

Bernoulli — closed form, icdf(q) = 0 if q ≤ 1 − p else 1.
Poisson, Binomial, NegativeBinomial — these have no closed-form quantile, and the special-function inverses (gammaincinv/betaincinv) don't apply since they invert the CDF in its continuous probability argument rather than over the integer support. This adds a shared discrete_icdf_via_search helper (in dist_math.py) that inverts the monotone logcdf by bisection, with geometric bracket expansion for the unbounded-support cases (Poisson, NegativeBinomial).

The helper uses fixed-length scan loops (no data-dependent until) so the graph remains convertible to the C, Numba, and JAX backends — verified for all three. As a side effect this enables pm.Truncated for these distributions, which falls back to icdf for inverse-transform sampling.

Tests compare against SciPy's ppf using the existing check_icdf helper.

A few other distributions in the same family (BetaBinomial, HyperGeometric) could reuse the same helper if you all want! This was very fun to work on. I met Daniel Lee at the Sloan conf back in March and have been using PyMc ever since--though I still enjoy Stan.

Related Issue

Related to Add icdf functions for distributions #6612

Checklist

Checked that the pre-commit linting/style checks pass
Included tests that prove the feature works
Added necessary documentation (docstrings)

Type of change

New feature / enhancement

read-the-docs-community · 2026-06-22T23:20:38Z

Documentation build overview

📚 pymc | 🛠️ Build #33565707 | 📁 Comparing 77e2f78 against latest (94787c0)

🔍 Preview build

165 files changed · ± 165 modified

± Modified

ricardoV94 · 2026-06-26T23:19:13Z

Can you split the bernoulli into it's own PR, it's much more trivial to review and approve than the ones requiring search

jacobpstein · 2026-06-29T20:37:22Z

Done! split the Bernoulli icdf into #8343. This PR now covers only the search-based distributions (Binomial, Poisson, NegativeBinomial).

ricardoV94 · 2026-07-02T15:07:45Z

+    whose quantile function has no closed form (e.g. ``Poisson``, ``Binomial``,
+    ``NegativeBinomial``).
+
+    The search uses fixed-length ``scan`` loops (rather than a data-dependent


Have to think about this, C/Numba are more than happy with data dependent loops.

Yeahhhh, the real constraint I was designing around is JAX (and static-shape lowering generally), which can't do a data-dependent number of steps.

I went with fixed-length scan so the same icdf graph works across all backends without special-casing, at the cost of always running max_iters iterations even after convergence. One option is to use a data-dependent while loop, accepting that JAX would be unsupported (or fall back) for these icdfs. Something to think about.

ricardoV94 · 2026-07-02T15:08:46Z

    return expr


+def discrete_icdf_via_search(logcdf, value, lower, upper=None, *, max_iters=64):


I may be misremembering, but one of the libraries I looked at also allowed some candidate bounds or something to inform the starting point of the search, in case the callers can offer something useful. Not saying we have to, but if there are obvious cases and it's not too much complexity would be good to discuss

I'm always open to learning some new libraries! Right now upper is already the big lever-- Binomial passes upper=n and skips the expansion phase. For the unbounded cases the obvious hint is the mean (mu for Poisson, n(1-p)/p for NegativeBinomial), so I could add an optional start that seeds the geometric expansion there instead of doubling from 1.

Payoff is modest though since bisection is already logarithmic, so a hint just trims a few iterations off phase 1. Low-complexity to add if you think it's worth it, otherwise happy to leave as a follow-up, and if there's an existing package whose approach we can reference here, even better

ricardoV94 · 2026-07-02T15:09:20Z

@jacobpstein implementation looks really clean. I don't have concrete feedback yet, have to marinate a bit on it. Let's some vague questions/thoughts above

The geometric bracket expansion can now be seeded with a caller-provided guess (typically the distribution mean) instead of always doubling from lower + 1. Poisson and NegativeBinomial pass their means; Binomial is unchanged since it already provides upper=n.

jacobpstein · 2026-07-13T13:35:25Z

Went ahead and added the start hint we talked about. Poisson seeds the bracket expansion with mu, NegativeBinomial with n(1-p)/p, and Binomial doesn't need it since upper=n already skips the expansion phase. Nice side effect: if the quantile you're asking for is at or below the guess (which it usually is), phase 1 is a no-op and we go straight to bisection.

On the loop question, I dug into this a bit more and unfortunately there's no middle ground. Adding an until condition to scan turns it into a while-scan, and JAX flat-out refuses those ("While Scan cannot yet be converted to JAX"), even with a fixed n_steps cap. So it really is fixed-length-everywhere vs. faster-loops-but-no-JAX. My vote is to keep the fixed-length version for now (each step is just one logcdf eval, and the start hint trims the effective work) and revisit if pytensor ever learns to lower while-scans to JAX. But if you'd rather have the speed and drop JAX for these, it's a two-line change, so happy to go either way.

Fun discovery while stress-testing at large mu: pytensor's C gammaincc (GammaQ in scalar/c_code/gamma.c) quietly falls apart for a ≈ x ≳ 1e5. For example, gammaincc(1e6, 1e6) comes back as 0.653 when scipy says 0.4999. The NR series and continued fraction both need O(√x) iterations in that regime and just give up at the MAXITER=1024 cap without complaining. Standalone graphs happen to dodge it via the scipy path, but inside scan the fused graph goes through C, so Poisson logcdf/icdf can drift at extreme mu on the C backend. That's totally independent of this PR, but worth knowing. I'll open a pytensor issue for it; the new large-mu test here sits at mu=1e4, comfortably below the danger zone.

codecov · 2026-07-13T13:46:36Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 90.12%. Comparing base (9d24260) to head (77e2f78).
⚠️ Report is 15 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #8339       +/-   ##
===========================================
+ Coverage   79.80%   90.12%   +10.32%     
===========================================
  Files         125      128        +3     
  Lines       20526    20733      +207     
===========================================
+ Hits        16380    18685     +2305     
+ Misses       4146     2048     -2098

Files with missing lines	Coverage Δ
pymc/distributions/discrete.py	`99.45% <100.00%> (+37.67%)`	⬆️
pymc/distributions/dist_math.py	`85.98% <100.00%> (+2.89%)`	⬆️
pymc/distributions/truncated.py	`99.48% <100.00%> (ø)`

... and 56 files with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

jacobpstein · 2026-07-13T13:47:33Z

Opened pymc-devs/pytensor#2291 for the gammaincc thing.

icdf graphs are always float-typed because out-of-domain values map to nan. When Truncated uses inverse-CDF sampling of a discrete base distribution, the resulting RV silently became float64, so e.g. Hurdle mixtures rejected it as a continuous component. Cast the icdf draw back to the base RV dtype.

jacobpstein marked this pull request as ready for review June 22, 2026 23:26

jacobpstein mentioned this pull request Jun 29, 2026

Implement icdf for Bernoulli #8343

Merged

Implement icdf for Poisson, Binomial, and NegativeBinomial

1fb6f6f

jacobpstein force-pushed the add-icdf-discrete-distributions branch from 1f662e1 to 1fb6f6f Compare June 29, 2026 20:37

jacobpstein changed the title ~~Implement icdf for Bernoulli, Poisson, Binomial, and NegativeBinomial~~ Implement icdf for Poisson, Binomial, and NegativeBinomial Jun 29, 2026

ricardoV94 reviewed Jul 2, 2026

View reviewed changes

jacobpstein added 2 commits July 13, 2026 10:03

Retrigger CI

77e2f78

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Implement icdf for Poisson, Binomial, and NegativeBinomial#8339

Implement icdf for Poisson, Binomial, and NegativeBinomial#8339
jacobpstein wants to merge 4 commits into
pymc-devs:mainfrom
jacobpstein:add-icdf-discrete-distributions

jacobpstein commented Jun 22, 2026 •

edited

Loading

Uh oh!

read-the-docs-community Bot commented Jun 22, 2026 •

edited

Loading

Uh oh!

ricardoV94 commented Jun 26, 2026

Uh oh!

jacobpstein commented Jun 29, 2026 •

edited

Loading

Uh oh!

ricardoV94 Jul 2, 2026

Uh oh!

jacobpstein Jul 2, 2026

Uh oh!

ricardoV94 Jul 2, 2026

Uh oh!

jacobpstein Jul 2, 2026

Uh oh!

ricardoV94 commented Jul 2, 2026

Uh oh!

jacobpstein commented Jul 13, 2026

Uh oh!

codecov Bot commented Jul 13, 2026 •

edited

Loading

Uh oh!

jacobpstein commented Jul 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		return expr


		def discrete_icdf_via_search(logcdf, value, lower, upper=None, *, max_iters=64):

Uh oh!

Conversation

jacobpstein commented Jun 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issue

Checklist

Type of change

Uh oh!

read-the-docs-community Bot commented Jun 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Documentation build overview

Uh oh!

ricardoV94 commented Jun 26, 2026

Uh oh!

jacobpstein commented Jun 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ricardoV94 Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

jacobpstein Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

jacobpstein Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

ricardoV94 commented Jul 2, 2026

Uh oh!

jacobpstein commented Jul 13, 2026

Uh oh!

codecov Bot commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

jacobpstein commented Jul 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

jacobpstein commented Jun 22, 2026 •

edited

Loading

read-the-docs-community Bot commented Jun 22, 2026 •

edited

Loading

jacobpstein commented Jun 29, 2026 •

edited

Loading

codecov Bot commented Jul 13, 2026 •

edited

Loading