From 83cb13865583fb1ec98d5ad5d4c0048f40e10e60 Mon Sep 17 00:00:00 2001 From: Juan Orduz Date: Fri, 24 Jul 2026 10:51:19 +0200 Subject: [PATCH 1/3] Use pytensor-distributions for Normal logp, logcdf, logccdf and icdf Co-Authored-By: Claude Fable 5 --- conda-envs/environment-alternative-backends.yml | 1 + conda-envs/environment-dev.yml | 1 + conda-envs/environment-docs.yml | 1 + conda-envs/environment-test.yml | 1 + conda-envs/windows-environment-dev.yml | 1 + conda-envs/windows-environment-test.yml | 1 + pymc/distributions/continuous.py | 12 +++++------- requirements-dev.txt | 1 + requirements.txt | 1 + 9 files changed, 13 insertions(+), 7 deletions(-) diff --git a/conda-envs/environment-alternative-backends.yml b/conda-envs/environment-alternative-backends.yml index 9b87fa7c7f..f8668d6de7 100644 --- a/conda-envs/environment-alternative-backends.yml +++ b/conda-envs/environment-alternative-backends.yml @@ -21,6 +21,7 @@ dependencies: - pandas>=0.24.0 - pip - pytensor>=3.2.2,<3.3 +- pytensor-distributions>=0.2.0 - python-graphviz - networkx - rich>=13.7.1 diff --git a/conda-envs/environment-dev.yml b/conda-envs/environment-dev.yml index 158ef1bfb5..69d9a06471 100644 --- a/conda-envs/environment-dev.yml +++ b/conda-envs/environment-dev.yml @@ -13,6 +13,7 @@ dependencies: - pandas>=0.24.0 - pip - pytensor>=3.2.2,<3.3 +- pytensor-distributions>=0.2.0 - python-graphviz - networkx - scipy>=1.4.1 diff --git a/conda-envs/environment-docs.yml b/conda-envs/environment-docs.yml index 10f8db8236..9a54a95df0 100644 --- a/conda-envs/environment-docs.yml +++ b/conda-envs/environment-docs.yml @@ -12,6 +12,7 @@ dependencies: - pandas>=0.24.0 - pip - pytensor>=3.2.2,<3.3 +- pytensor-distributions>=0.2.0 - python-graphviz - rich>=13.7.1 - scipy>=1.4.1 diff --git a/conda-envs/environment-test.yml b/conda-envs/environment-test.yml index 56a71ae4ca..297709b665 100644 --- a/conda-envs/environment-test.yml +++ b/conda-envs/environment-test.yml @@ -16,6 +16,7 @@ dependencies: - pip - polyagamma - pytensor>=3.2.2,<3.3 +- pytensor-distributions>=0.2.0 - python-graphviz - networkx - rich>=13.7.1 diff --git a/conda-envs/windows-environment-dev.yml b/conda-envs/windows-environment-dev.yml index e0e4d1adb3..ae2083320a 100644 --- a/conda-envs/windows-environment-dev.yml +++ b/conda-envs/windows-environment-dev.yml @@ -13,6 +13,7 @@ dependencies: - pandas>=0.24.0 - pip - pytensor>=3.2.2,<3.3 +- pytensor-distributions>=0.2.0 - python-graphviz - networkx - rich>=13.7.1 diff --git a/conda-envs/windows-environment-test.yml b/conda-envs/windows-environment-test.yml index 221244a916..08a94bd1e0 100644 --- a/conda-envs/windows-environment-test.yml +++ b/conda-envs/windows-environment-test.yml @@ -17,6 +17,7 @@ dependencies: - pip - polyagamma - pytensor>=3.2.2,<3.3 +- pytensor-distributions>=0.2.0 - python-graphviz - networkx - rich>=13.7.1 diff --git a/pymc/distributions/continuous.py b/pymc/distributions/continuous.py index 4b3eeb53d0..23f8b1d2bf 100644 --- a/pymc/distributions/continuous.py +++ b/pymc/distributions/continuous.py @@ -51,6 +51,7 @@ from pytensor.tensor.random.op import RandomVariable from pytensor.tensor.random.utils import normalize_size_param from pytensor.tensor.variable import TensorConstant, TensorVariable +from pytensor_distributions import normal as ptd_normal from pymc.distributions.custom import CustomDist from pymc.logprob.abstract import _logprob_helper @@ -524,32 +525,29 @@ def support_point(rv, size, mu, sigma): return mu def logp(value, mu, sigma): - res = -0.5 * pt.pow((value - mu) / sigma, 2) - pt.log(pt.sqrt(2.0 * np.pi)) - pt.log(sigma) return check_parameters( - res, + ptd_normal.logpdf(value, mu, sigma), sigma > 0, msg="sigma > 0", ) def logcdf(value, mu, sigma): return check_parameters( - normal_lcdf(mu, sigma, value), + ptd_normal.logcdf(value, mu, sigma), sigma > 0, msg="sigma > 0", ) def logccdf(value, mu, sigma): return check_parameters( - normal_lccdf(mu, sigma, value), + ptd_normal.logsf(value, mu, sigma), sigma > 0, msg="sigma > 0", ) def icdf(value, mu, sigma): - res = mu + sigma * -np.sqrt(2.0) * pt.erfcinv(2 * value) - res = check_icdf_value(res, value) return check_icdf_parameters( - res, + ptd_normal.ppf(value, mu, sigma), sigma > 0, msg="sigma > 0", ) diff --git a/requirements-dev.txt b/requirements-dev.txt index 8a10c9b0d3..374083ff9a 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -17,6 +17,7 @@ pandas>=0.24.0 polyagamma pre-commit>=2.8.0 pymc-sphinx-theme>=0.16.0 +pytensor-distributions>=0.2.0 pytensor>=3.2.2,<3.3 pytest-cov>=2.5 pytest>=3.0 diff --git a/requirements.txt b/requirements.txt index 52ca0d21dc..b178db602e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,6 +4,7 @@ cloudpickle numpy>=1.25.0 pandas>=0.24.0 pytensor>=3.2.2,<3.3 +pytensor-distributions>=0.2.0 rich>=13.7.1 scipy>=1.4.1 threadpoolctl>=3.1.0,<4.0.0 From 69a8df407599885e173de4a143efef659f1b6010 Mon Sep 17 00:00:00 2001 From: Juan Orduz Date: Fri, 7 Aug 2026 17:52:41 +0200 Subject: [PATCH 2/3] Pin pytensor-distributions at <=0.3.0 Per review, cap the new dependency so a future breaking release does not force a rushed fix. Co-Authored-By: Claude Fable 5 --- conda-envs/environment-alternative-backends.yml | 2 +- conda-envs/environment-dev.yml | 2 +- conda-envs/environment-docs.yml | 2 +- conda-envs/environment-test.yml | 2 +- conda-envs/windows-environment-dev.yml | 2 +- conda-envs/windows-environment-test.yml | 2 +- requirements-dev.txt | 2 +- requirements.txt | 2 +- 8 files changed, 8 insertions(+), 8 deletions(-) diff --git a/conda-envs/environment-alternative-backends.yml b/conda-envs/environment-alternative-backends.yml index f8668d6de7..375f00d692 100644 --- a/conda-envs/environment-alternative-backends.yml +++ b/conda-envs/environment-alternative-backends.yml @@ -21,7 +21,7 @@ dependencies: - pandas>=0.24.0 - pip - pytensor>=3.2.2,<3.3 -- pytensor-distributions>=0.2.0 +- pytensor-distributions>=0.2.0,<=0.3.0 - python-graphviz - networkx - rich>=13.7.1 diff --git a/conda-envs/environment-dev.yml b/conda-envs/environment-dev.yml index 69d9a06471..d800e247fa 100644 --- a/conda-envs/environment-dev.yml +++ b/conda-envs/environment-dev.yml @@ -13,7 +13,7 @@ dependencies: - pandas>=0.24.0 - pip - pytensor>=3.2.2,<3.3 -- pytensor-distributions>=0.2.0 +- pytensor-distributions>=0.2.0,<=0.3.0 - python-graphviz - networkx - scipy>=1.4.1 diff --git a/conda-envs/environment-docs.yml b/conda-envs/environment-docs.yml index 9a54a95df0..25b9484073 100644 --- a/conda-envs/environment-docs.yml +++ b/conda-envs/environment-docs.yml @@ -12,7 +12,7 @@ dependencies: - pandas>=0.24.0 - pip - pytensor>=3.2.2,<3.3 -- pytensor-distributions>=0.2.0 +- pytensor-distributions>=0.2.0,<=0.3.0 - python-graphviz - rich>=13.7.1 - scipy>=1.4.1 diff --git a/conda-envs/environment-test.yml b/conda-envs/environment-test.yml index 297709b665..e833ba3437 100644 --- a/conda-envs/environment-test.yml +++ b/conda-envs/environment-test.yml @@ -16,7 +16,7 @@ dependencies: - pip - polyagamma - pytensor>=3.2.2,<3.3 -- pytensor-distributions>=0.2.0 +- pytensor-distributions>=0.2.0,<=0.3.0 - python-graphviz - networkx - rich>=13.7.1 diff --git a/conda-envs/windows-environment-dev.yml b/conda-envs/windows-environment-dev.yml index ae2083320a..ad6b1ef9e5 100644 --- a/conda-envs/windows-environment-dev.yml +++ b/conda-envs/windows-environment-dev.yml @@ -13,7 +13,7 @@ dependencies: - pandas>=0.24.0 - pip - pytensor>=3.2.2,<3.3 -- pytensor-distributions>=0.2.0 +- pytensor-distributions>=0.2.0,<=0.3.0 - python-graphviz - networkx - rich>=13.7.1 diff --git a/conda-envs/windows-environment-test.yml b/conda-envs/windows-environment-test.yml index 08a94bd1e0..8d9b5538af 100644 --- a/conda-envs/windows-environment-test.yml +++ b/conda-envs/windows-environment-test.yml @@ -17,7 +17,7 @@ dependencies: - pip - polyagamma - pytensor>=3.2.2,<3.3 -- pytensor-distributions>=0.2.0 +- pytensor-distributions>=0.2.0,<=0.3.0 - python-graphviz - networkx - rich>=13.7.1 diff --git a/requirements-dev.txt b/requirements-dev.txt index 374083ff9a..ab615a17cc 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -17,7 +17,7 @@ pandas>=0.24.0 polyagamma pre-commit>=2.8.0 pymc-sphinx-theme>=0.16.0 -pytensor-distributions>=0.2.0 +pytensor-distributions>=0.2.0,<=0.3.0 pytensor>=3.2.2,<3.3 pytest-cov>=2.5 pytest>=3.0 diff --git a/requirements.txt b/requirements.txt index b178db602e..51f984c3c5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,7 +4,7 @@ cloudpickle numpy>=1.25.0 pandas>=0.24.0 pytensor>=3.2.2,<3.3 -pytensor-distributions>=0.2.0 +pytensor-distributions>=0.2.0,<=0.3.0 rich>=13.7.1 scipy>=1.4.1 threadpoolctl>=3.1.0,<4.0.0 From 6ca7f7c4af9e056ffdcbe1ae896efced58d0810f Mon Sep 17 00:00:00 2001 From: Juan Orduz Date: Fri, 7 Aug 2026 17:52:42 +0200 Subject: [PATCH 3/3] Add one-off logp_dlogp benchmark for Normal (to be removed before merge) Before/after analysis of model.logp_dlogp_function(ravel_inputs=True) requested in review: dprint with print_memory_map=True, timings, op counts, and value checks. See PTD_NORMAL_BENCHMARK.md for results. Co-Authored-By: Claude Fable 5 --- PTD_NORMAL_BENCHMARK.md | 247 ++++++++ benchmark_normal_ptd.py | 221 +++++++ ptd_benchmark_results/after/icdf_dprint.txt | 40 ++ ptd_benchmark_results/after/meta.json | 11 + .../normal_censored_logp_dlogp_dprint.txt | 571 ++++++++++++++++++ .../after/normal_logp_dlogp_dprint.txt | 141 +++++ ptd_benchmark_results/after/op_counts.json | 19 + ptd_benchmark_results/after/timings.json | 20 + ptd_benchmark_results/after/values.json | 28 + ptd_benchmark_results/before/icdf_dprint.txt | 24 + ptd_benchmark_results/before/meta.json | 11 + .../normal_censored_logp_dlogp_dprint.txt | 500 +++++++++++++++ .../before/normal_logp_dlogp_dprint.txt | 141 +++++ ptd_benchmark_results/before/op_counts.json | 19 + ptd_benchmark_results/before/timings.json | 20 + ptd_benchmark_results/before/values.json | 28 + 16 files changed, 2041 insertions(+) create mode 100644 PTD_NORMAL_BENCHMARK.md create mode 100644 benchmark_normal_ptd.py create mode 100644 ptd_benchmark_results/after/icdf_dprint.txt create mode 100644 ptd_benchmark_results/after/meta.json create mode 100644 ptd_benchmark_results/after/normal_censored_logp_dlogp_dprint.txt create mode 100644 ptd_benchmark_results/after/normal_logp_dlogp_dprint.txt create mode 100644 ptd_benchmark_results/after/op_counts.json create mode 100644 ptd_benchmark_results/after/timings.json create mode 100644 ptd_benchmark_results/after/values.json create mode 100644 ptd_benchmark_results/before/icdf_dprint.txt create mode 100644 ptd_benchmark_results/before/meta.json create mode 100644 ptd_benchmark_results/before/normal_censored_logp_dlogp_dprint.txt create mode 100644 ptd_benchmark_results/before/normal_logp_dlogp_dprint.txt create mode 100644 ptd_benchmark_results/before/op_counts.json create mode 100644 ptd_benchmark_results/before/timings.json create mode 100644 ptd_benchmark_results/before/values.json diff --git a/PTD_NORMAL_BENCHMARK.md b/PTD_NORMAL_BENCHMARK.md new file mode 100644 index 0000000000..17913d1ef1 --- /dev/null +++ b/PTD_NORMAL_BENCHMARK.md @@ -0,0 +1,247 @@ +# One-off `logp_dlogp` analysis for #8376 (to be removed before merge) + +Requested in review: an analysis of `model.logp_dlogp_function(ravel_inputs=True)` — +the compiled graph (`print_memory_map=True`) and a timing that confirms no +graph/eval regression after delegating Normal's `logp`/`logcdf`/`logccdf`/`icdf` +to [pytensor-distributions](https://github.com/pymc-devs/pytensor-distributions). + +Produced by `benchmark_normal_ptd.py` (repo root). Raw artifacts (dprints, +timings, op counts, values) live in `ptd_benchmark_results/{before,after}/`. + +## Reproduce + +```bash +python benchmark_normal_ptd.py run --label after # on this branch +git switch --detach $(git merge-base upstream/main HEAD) +python benchmark_normal_ptd.py run --label before # old implementation +git switch normal-pytensor-distributions +python benchmark_normal_ptd.py compare +``` + +## Environment + +| | | +|---|---| +| before / after commits | `3b661c7e5` (merge-base with main) / `83cb13865` (this PR) | +| pytensor / pytensor-distributions | 3.2.2 / 0.2.0 | +| numpy / python | 2.4.6 / 3.13.11 | +| machine / floatX | macOS 26.5.2 arm64 (M-series) / float64 | + +## Benchmarks + +1. **`normal`** — `mu ~ Normal`, `sigma ~ HalfNormal`, `obs ~ Normal(mu, sigma)` with + N=1000 observed. Exercises `Normal.logp` (scalar prior + vector likelihood). +2. **`normal_censored`** — same plus `Censored(Normal.dist(mu, sigma), -1, 1)` with + N=500 observed, which pulls `Normal.logcdf` **and** `Normal.logccdf` into the + logp graph. +3. **`icdf`** — standalone compiled `pm.icdf(pm.Normal.dist(1.5, 2.5), q)` (not part + of any logp graph, but also changed by this PR). + +Timings: `timeit.repeat(repeat=30, number=200)`, best (min) per call. + +## Results + +| bench | graph | before min | after min | ratio | values | +|---|---|---:|---:|---:|---| +| `normal` logp_dlogp | **byte-identical dprint** (141 lines) | 2.14 µs | 2.17 µs | **1.02** (noise) | bit-identical | +| `normal_censored` logp_dlogp | differs (see below) | 26.9 µs | 36.1 µs | **1.34** | logp bit-identical; dlogp[0] differs by 1 ulp | +| `icdf` | differs by design (see below) | 7.9 µs | 8.8 µs | 1.12 | ≤ 2 ulp | + +### `normal`: no regression ✅ + +The delegated `logpdf` is the same expression as the old inline formula, so after +rewrites the compiled `logp_dlogp` graph is **byte-for-byte identical** and evaluation +time is unchanged. This is the core code path for the vast majority of models. + +### `normal_censored`: graph + eval regression from `logsf` ⚠️ + +`pytensor_distributions.normal` implements + +```python +def logsf(x, mu, sigma): + return logcdf(-x, -mu, sigma) +``` + +which is mathematically identical to the old `normal_lccdf`, but graph-wise it is +not: `-x` bakes a **negated copy of the observed data into the graph as a second +constant** (visible in the dprint as `[-5.244064 ... 58192e-01]`, duplicating the +censored data's memory), and the standardized value inside becomes +`(-x - (-mu))/sigma`, which the rewriter does not recombine into `-((x - mu)/sigma)` +— so `logcdf` and `logccdf` no longer share the `z = (x - mu)/sigma` subexpression. +Consequences on this model: + +- `DimShuffle` count 8 → 12; the censored fused elemwise grows from 5 to 9 + `reduce[add]` outputs; dprint grows 500 → 571 lines. +- eval time 26.9 µs → 36.1 µs (**1.34×**) for N=500 censored points. +- `dlogp` wrt `mu` differs in the last ulp (different but equivalent expression + arrangement); logp itself is bit-identical. + +The old PyMC `normal_lccdf` applies the erfcx switch directly to `z` +(`switch(z > 1, log(erfcx(z/√2)/2) − z²/2, log1p(−erfc(−z/√2)/2))`). Writing +`logsf` that way in pytensor-distributions would restore subexpression sharing and +remove the duplicated constant — a small upstream fix; alternatively PyMC can keep +its own `logccdf` until then. + +### `icdf`: intentional behavior difference, minor micro-overhead + +- Old: `check_icdf_value` wraps `−erfcinv(2q)·σ√2 + μ` in a single + `switch(0 ≤ q ≤ 1, ..., nan)`. +- New: the package's `ppf` uses `erfinv(2q − 1)` inside `ppf_bounds_cont`, which + adds explicit branches `q<0 or q>1 → nan`, `q=0 → −inf`, `q=1 → +inf`. Same + results (boundary infinities previously arose from `erfcinv` itself), values + match to ≤ 2 ulp (`erfinv` vs `erfcinv`), and the extra scalar switches cost + ~0.9 µs on a 1000-point batch. Not part of logp graphs; irrelevant for sampling. + +## Compiled graph (`print_memory_map=True`), after — `normal` model + +
ptd_benchmark_results/after/normal_logp_dlogp_dprint.txt (byte-identical to before) + +```text +Composite{...}.0 [id A] d={0: [1]} 8 + ├─ Subtensor{i} [id B] v={0: [0]} 1 + │ ├─ joined_inputs [id C] + │ └─ 1 [id D] + ├─ FusedElemwise{Composite{...}, reduce[add@(0,), add@(0,), add@(0,)]}.0 [id E] 6 + │ ├─ Composite{...}.1 [id F] 5 + │ │ └─ ExpandDims{axis=0} [id G] v={0: [0]} 4 + │ │ └─ Exp [id H] 3 + │ │ └─ Subtensor{i} [id B] v={0: [0]} 1 + │ │ └─ ··· + │ ├─ Composite{...}.0 [id F] 5 + │ │ └─ ··· + │ ├─ ExpandDims{axis=0} [id G] v={0: [0]} 4 + │ │ └─ ··· + │ ├─ obs{[ 1.320334 ... 56013e+00]} [id I] + │ └─ Subtensor{:stop} [id J] v={0: [0]} 0 + │ ├─ joined_inputs [id C] + │ └─ 1 [id K] + ├─ Subtensor{i} [id L] v={0: [0]} 2 + │ ├─ joined_inputs [id C] + │ └─ 0 [id M] + └─ Exp [id H] 3 + └─ ··· +Join{axis=0} [id N] 12 + ├─ Composite{...}.0 [id O] d={0: [0], 1: [1]} 11 + │ ├─ ExpandDims{axis=0} [id G] v={0: [0]} 4 + │ │ └─ ··· + │ ├─ ExpandDims{axis=0} [id P] v={0: [0]} 7 + │ │ └─ FusedElemwise{Composite{...}, reduce[add@(0,), add@(0,), add@(0,)]}.1 [id E] 6 + │ │ └─ ··· + │ ├─ Subtensor{:stop} [id J] v={0: [0]} 0 + │ │ └─ ··· + │ ├─ ExpandDims{axis=0} [id Q] v={0: [0]} 10 + │ │ └─ Composite{...}.1 [id A] d={0: [1]} 8 + │ │ └─ ··· + │ └─ ExpandDims{axis=0} [id R] v={0: [0]} 9 + │ └─ FusedElemwise{Composite{...}, reduce[add@(0,), add@(0,), add@(0,)]}.2 [id E] 6 + │ └─ ··· + └─ Composite{...}.1 [id O] d={0: [0], 1: [1]} 11 + └─ ··· + +Inner graphs: + +Composite{...} [id A] d={0: [1]} + ← add [id S] + ├─ -3.221523658174155 [id T] + ├─ mul [id U] + │ ├─ -0.5 [id V] + │ └─ sqr [id W] + │ └─ mul [id X] + │ ├─ 0.1 [id Y] + │ └─ i2 [id Z] + ├─ Switch [id BA] + │ ├─ GE [id BB] + │ │ ├─ i3 [id BC] + │ │ └─ 0.0 [id BD] + │ ├─ add [id BE] + │ │ ├─ -1.83522929514961 [id BF] + │ │ └─ mul [id BG] + │ │ ├─ -0.5 [id V] + │ │ └─ sqr [id BH] + │ │ └─ mul [id BI] + │ │ ├─ 0.2 [id BJ] + │ │ └─ i3 [id BC] + │ └─ -inf [id BK] + ├─ i0 [id BL] + └─ i1 [id BM] + ← GE [id BB] + └─ ··· + +FusedElemwise{Composite{...}, reduce[add@(0,), add@(0,), add@(0,)]} [id E] + ← Sum{axes=None} [id BN] + └─ Composite{...}.0 [id BO] + ├─ i0 [id BP] + ├─ i1 [id BQ] + ├─ i2 [id BR] + ├─ i3 [id BS] + └─ i4 [id BT] + ← Sum{axes=None} [id BU] + └─ Composite{...}.1 [id BO] + └─ ··· + ← Sum{axes=None} [id BV] + └─ Composite{...}.2 [id BO] + └─ ··· + +Composite{...} [id F] + ← log [id BW] + └─ i0 [id BL] + ← GT [id BX] + ├─ i0 [id BL] + └─ 0 [id BY] + +Composite{...} [id O] d={0: [0], 1: [1]} + ← add [id BZ] + ├─ mul [id CA] + │ ├─ -0.0100000 ... 0000000002 [id CB] + │ └─ i2 [id Z] + └─ true_div [id CC] + ├─ true_div [id CD] + │ ├─ i1 [id BM] + │ └─ i0 [id BL] + └─ i0 [id BL] + ← add [id CE] + ├─ -999.0 [id CF] + ├─ Switch [id CG] + │ ├─ i3 [id CH] + │ ├─ mul [id CI] + │ │ ├─ -0.04000000000000001 [id CJ] + │ │ ├─ i0 [id BL] + │ │ └─ i0 [id BL] + │ └─ 0.0 [id CK] + └─ mul [id CL] + ├─ true_div [id CM] + │ ├─ i4 [id CN] + │ └─ sqr [id CO] + │ └─ i0 [id BL] + └─ i0 [id BL] + +Composite{...} [id BO] + ← Switch [id CP] + ├─ i0 [id CQ] + ├─ sub [id CR] + │ ├─ add [id CS] + │ │ ├─ -0.9189385332046727 [id CT] + │ │ └─ mul [id CU] + │ │ ├─ -0.5 [id CV] + │ │ └─ sqr [id CW] + │ │ └─ true_div [id CX] + │ │ ├─ sub [id CY] + │ │ │ ├─ i3 [id BC] + │ │ │ └─ i4 [id CN] + │ │ └─ i2 [id Z] + │ └─ i1 [id BM] + └─ -inf [id CZ] + ← sub [id CY] + └─ ··· + ← mul [id DA] + ├─ true_div [id CX] + │ └─ ··· + └─ sub [id CY] + └─ ··· +``` + +
+ +Full dprints for all benchmarks (before and after) are committed under +`ptd_benchmark_results/` — `diff` the `before/` and `after/` files to see the +graph changes described above. diff --git a/benchmark_normal_ptd.py b/benchmark_normal_ptd.py new file mode 100644 index 0000000000..c3ae27dd32 --- /dev/null +++ b/benchmark_normal_ptd.py @@ -0,0 +1,221 @@ +# Copyright 2026 - present The PyMC Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""One-off benchmark for pymc-devs/pymc#8376 (to be removed before merge). + +Compares the compiled ``model.logp_dlogp_function(ravel_inputs=True)`` graph and +its evaluation time before and after delegating Normal's logp/logcdf/logccdf/icdf +to pytensor-distributions, as requested in review. + +Usage (run twice, then compare):: + + python benchmark_normal_ptd.py run --label after # on the PR branch + git switch --detach $(git merge-base upstream/main HEAD) + python benchmark_normal_ptd.py run --label before # on main (old impl) + git switch normal-pytensor-distributions + python benchmark_normal_ptd.py compare + +Outputs go to ``ptd_benchmark_results/