Add pmd.CustomDist — dims-aware custom distribution for pymc.dims by williambdean · Pull Request #8311 · pymc-devs/pymc

williambdean · 2026-05-20T12:18:06Z

Adds CustomDist to pymc.dims.distributions, a sibling to pm.CustomDist that operates on XTensorVariable with named dims.

Two construction paths:

Symbolic (dist= kwarg): receives XTensorVariable params, returns an XTensorVariable RV (e.g., composing pmd.Normal.dist). Auto-derives logp from inner XRV nodes.
Black-box (logp= kwarg): dynamically creates a RandomVariable subclass; dispatches _logprob, _logcdf, _support_point. The value arrives as XTensorVariable; use .values for pt.* ops or ptx.* for dim-aware ops.

Key design points:

Params go through the same DimDistribution._as_xtensor path as pmd.Normal etc. — identical behavior (scalars auto-convert, non-scalars require dims).
User callables (logp, logcdf, support_point) captured in closures to avoid Python descriptor protocol issues.
Dynamic RandomVariable subclass sets only signature (not ndim_supp/ndims_params) to avoid deprecation warnings.

ricardoV94 · 2026-05-20T12:22:45Z

Black-box (logp= kwarg),

These are orthogonal to having a dist argument. You can have dist with logp (or without, maybe it derives it). The only incompatible case is dist AND random, since they both represent the random path

ricardoV94 · 2026-05-20T12:26:38Z

+    return func
+
+
+def _default_support_point(rv, size, *rv_inputs, rv_name=None, has_fallback=False):


I think this is just rv.zeros_like() ?

no ones_like or zeros_like on xtensor FYI

import pytensor.xtensor as px x = px.xtensor("x", dims=("covariate", )) x.ones_like()

--------------------------------------------------------------------------- AttributeError Traceback (most recent call last) Cell In[35], line 1 ----> 1 px.xtensor("x", dims=("covariate", )).ones_like() AttributeError: 'XTensorVariable' object has no attribute 'ones_like'

but can use px.zeros_like

codecov · 2026-05-20T12:32:51Z

Codecov Report

❌ Patch coverage is 86.60287% with 28 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.85%. Comparing base (971800b) to head (c570454).

Files with missing lines	Patch %	Lines
pymc/dims/distributions/custom.py	84.69%	28 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (971800b) and HEAD (c570454). Click for more details.

HEAD has 8 uploads less than BASE

Flag BASE (971800b) HEAD (c570454)

32 24

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #8311       +/-   ##
===========================================
- Coverage   91.72%   79.85%   -11.88%     
===========================================
  Files         125      126        +1     
  Lines       20526    20716      +190     
===========================================
- Hits        18828    16542     -2286     
- Misses       1698     4174     +2476

Files with missing lines	Coverage Δ
pymc/dims/distributions/__init__.py	`100.00% <100.00%> (ø)`
pymc/dims/distributions/core.py	`91.83% <100.00%> (+0.12%)`	⬆️
pymc/distributions/custom.py	`78.08% <100.00%> (-17.48%)`	⬇️
pymc/dims/distributions/custom.py	`84.69% <84.69%> (ø)`

... and 41 files with indirect coverage changes

🚀 New features to boost your workflow:

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

ricardoV94 · 2026-05-20T12:34:04Z

Good start, I think we should drop the random argument and a lot of complexity falls out of the way

read-the-docs-community · 2026-05-20T12:34:59Z

Documentation build overview

📚 pymc | 🛠️ Build #32982419 | 📁 Comparing c570454 against latest (f09f6b6)

🔍 Preview build

5 files changed · ± 5 modified

± Modified

Supports both symbolic (dist=) and black-box (logp=) paths, enabling user-defined distributions with named dims. The symbolic path auto-derives logprob from inner XRV nodes; the black-box path creates a dynamic RandomVariable subclass and registers _logprob dispatches that reconstruct XTensorVariables for the value and dims-bearing params.

Covers both symbolic (dist=) and black-box (logp=/random=) paths: graph comparison against regular distributions, dim propagation, observed data, custom support points, and model variables as params.

…signature inference, fix compound dists - Replace compiled-function + graph-walking hybrid path with DimSymbolicRandomVariable(SymbolicRandomVariable) + OpFromGraph - Deduplicate _infer_dims_signature / _infer_final_signature - Add XElemwise support to expand_dist_dims for compound dists - Drop _forward_dim_lengths, enforce strict XTensorVariable output - Add tests: compound non-XRV output, hybrid support_point