Use pytensor-distributions for Normal logp, logcdf, logccdf and icdf

[juanitorduz]

Description

Proof of concept for delegating distribution math to the new pytensor-distributions package, starting with Normal.

• Normal.logp, logcdf, logccdf and icdf now call pytensor_distributions.normal.logpdf / logcdf / logsf / ppf instead of the inline formulas. The expressions are mathematically identical to the previous ones (logsf uses the same erfcx-based asymptotic switch as normal_lccdf, and erfinv(2q-1) ≡ -erfcinv(2q) for the icdf), so numerical behavior is unchanged.
• Parameter validation stays in PyMC: the package does no validation, so the check_parameters / check_icdf_parameters (sigma > 0) wrappers are kept. The check_icdf_value wrapper was dropped because ppf already bounds q (NaN outside [0, 1], ±inf at the edges), matching the previous behavior.
• dist_math.normal_lcdf / normal_lccdf are untouched — they are still used by TruncatedNormal, LogNormal, ExGaussian, OrderedProbit, etc.
• Added pytensor-distributions>=0.2.0 (available on PyPI and conda-forge, noarch) to requirements.txt and all conda env files; requirements-dev.txt regenerated by the pip-from-conda hook.

Existing scipy-comparison tests (TestMatchesScipy::test_normal, tests/logprob/test_abstract.py including the logccdf graph and numerical-stability checks) cover the change; no new tests needed.

Related Issue

• Closes #
• Related to https://github.com/pymc-devs/pytensor-distributions

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works (covered by existing tests)
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

🤖 Generated with Claude Code

[ricardoV94]

I need one analysis of the logp_dlogp that model.logp_dlogp_function(ravel=True) does, both the graph (print_memory_map=True), and a timing that confirms no graph/eval regression. This won't be in the CI it's a one time thing for every dist we introduce

[read-the-docs-community]

Documentation build overview

📚 pymc | 🛠️ Build #33962795 | 📁 Comparing 6ca7f7c against latest (3b661c7)

  🔍 Preview build  

7 files changed · + 1 added · ± 6 modified

+ Added

• _modules/pytensor/tensor/special.html

± Modified

• glossary.html
• _modules/index.html
• api/math.html
• _modules/pymc/distributions/continuous.html
• _modules/pytensor/tensor/math.html
• api/generated/classmethods/pymc.backends.zarr.ZarrTrace.to_inferencedata.html

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 91.85%. Comparing base (3b661c7) to head (6ca7f7c).
⚠️ Report is 12 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #8376      +/-   ##
==========================================
+ Coverage   91.82%   91.85%   +0.02%     
==========================================
  Files         128      128              
  Lines       21110    21257     +147     
==========================================
+ Hits        19384    19525     +141     
- Misses       1726     1732       +6     

Files with missing lines	Coverage Δ
pymc/distributions/continuous.py	98.05% <100.00%> (-0.01%)	⬇️

... and 15 files with indirect coverage changes

🚀 New features to boost your workflow:

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

[juanitorduz]

I need one analysis of the logp_dlogp that model.logp_dlogp_function(ravel=True) does, both the graph (print_memory_map=True), and a timing that confirms no graph/eval regression. This won't be in the CI it's a one time thing for every dist we introduce

Makes sense @ricardoV94 , will you do it or do you want me to put together a little script?

[ricardoV94]

I need one analysis of the logp_dlogp that model.logp_dlogp_function(ravel=True) does, both the graph (print_memory_map=True), and a timing that confirms no graph/eval regression. This won't be in the CI it's a one time thing for every dist we introduce

Makes sense @ricardoV94 , will you do it or do you want me to put together a little script?

if you have capacity now go ahead, I'm swamped until next week

[juanitorduz]

I need one analysis of the logp_dlogp that model.logp_dlogp_function(ravel=True) does, both the graph (print_memory_map=True), and a timing that confirms no graph/eval regression. This won't be in the CI it's a one time thing for every dist we introduce

@ricardoV94 in 6ca7f7c I asked Claude tu run this one-off analysis. Is this helpful? You can iterate over it. I will remove it if we decide to merge

[juanitorduz]

@ricardoV94 here is the logp_dlogp_function(ravel_inputs=True) analysis (script benchmark_normal_ptd.py + results in PTD_NORMAL_BENCHMARK.md and raw dprints/timings under ptd_benchmark_results/ — all to be removed before merge). Before = merge-base 3b661c7, after = this PR; pytensor 3.2.2, pytensor-distributions 0.2.0, float64, macOS arm64.

bench	graph (print_memory_map=True dprint)	before min	after min	ratio	values
Normal prior + N=1000 likelihood	byte-identical	2.14 µs	2.17 µs	1.02 (noise)	bit-identical
same + Censored(Normal) N=500	differs	26.9 µs	36.1 µs	1.34	logp bit-identical, dlogp 1 ulp
standalone icdf (1000 q's)	differs (by design)	7.9 µs	8.8 µs	1.12	≤ 2 ulp

logp: no regression. The package's logpdf/logcdf are the same expressions as the old inline formulas, so the compiled graph is byte-for-byte identical and timing is unchanged. ✅

logccdf: real graph/eval regression, isolated to models that use it (e.g. censoring). The package implements logsf(x, mu, sigma) = logcdf(-x, -mu, sigma), which is mathematically equal to the old normal_lccdf but graph-wise (a) bakes a negated copy of the observed data into the graph as a second constant, and (b) turns the standardized value into (-x - (-mu))/sigma, which the rewriter doesn't recombine into -z — so logcdf/logccdf stop sharing the z = (x - mu)/sigma subexpression. Net effect on the censored model: DimShuffle 8→12, the censored fused elemwise grows from 5 to 9 reduce[add] outputs, 1.34× eval time. Writing logsf in pytensor-distributions with the erfcx switch applied directly to z (as normal_lccdf does) should restore graph identity; alternatively this PR can keep PyMC's logccdf until that lands upstream. Happy to open the upstream issue/PR.

icdf: intentional difference, cosmetic cost. The package's ppf uses erfinv(2q−1) inside ppf_bounds_cont, which adds explicit q=0→-inf / q=1→+inf / out-of-bounds→nan switches (previously the infinities came from erfcinv itself and check_icdf_value handled nan). Values agree to ≤2 ulp; not part of logp graphs.

Also applied the <=0.3.0 pin across requirements.txt and all conda env files.

🤖 Generated with Claude Code

[ricardoV94]

Yeha that's the sort of stuff. It should check symbolic mu/log(sigma) and variants with constant and symbolic data. data is not "always" constant.

The logsf should be patched upstream yes