Add WeightedZeroSumNormal distribution and transform

[PabloRoque]

Description

Adds pm.WeightedZeroSumNormal (and pmd.WeightedZeroSumNormal in the dims API): a generalization of ZeroSumNormal where draws satisfy a weighted zero-sum constraint sum(weights * x) = 0 along the last axis, i.e.

$$WZSN(\sigma, w) = N\big(0,; \sigma^2 (I_n - u u^\top)\big), \qquad u = w / \lVert w \rVert$$

With equal weights it reproduces ZeroSumNormal exactly — pointwise transform parity and logp parity are covered by tests.

Motivation. In hierarchical models with lopsided group sizes, centring group deviations under a weighted mean instead of the uniform mean decorrelates the shared mean from the dominant unit's deviation. The motivating case is MMM campaign hierarchies (e.g. a 90/5/3/2 spend split within a channel), where the likelihood pins down the spend-weighted combination of effects.

Implementation notes.

• New WeightedZeroSumTransform (tensor and dims variants): the restriction of the Householder reflection sending u to -e_n. The v = u + e_n sign choice keeps the working denominator at 1 + u_n ≥ 1 for positive weights, so there is no near-singular regime. The map is an isometry, hence log_jac_det = 0.
• New internal WeightedZeroSumNormalRV (SymbolicRandomVariable) rather than extending ZeroSumNormalRV, keeping the existing RV's input signature untouched.
• The RV signature gives sigma a scalar core, "[rng],[size],(),(n)->[rng],(n)"; the dims xrv_op therefore needs no dummy support-shape dim.
• Weights are cast to floatX at every entry point, and the logp's on-constraint check uses a dtype-aware tolerance (float32 draws satisfy the constraint only to ~1e-7).
• The dims transform accepts symbolic weights (e.g. pmd.Data), so set_data keeps the constraint in sync.
• Scope: a single constrained axis (the last). Multi-axis with per-axis weight vectors is a possible follow-up if there is interest.

Happy to adjust the API (e.g. folding this into ZeroSumNormal as a weights= kwarg was considered and rejected to keep the existing RV signature and class behavior untouched — but open to maintainer preference).

Related Issue

• Closes #
• Related to Implement more Dims distributions #7874

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• 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

[welcome]

]
💖 Thanks for opening this pull request! 💖 The PyMC community really appreciates your time and effort to contribute to the project. Please make sure you have read our Contributing Guidelines and filled in our pull request template to the best of your ability.

[read-the-docs-community]

Documentation build overview

📚 pymc | 🛠️ Build #34222579 | 📁 Comparing d394c9b against latest (da8fc47)

  🔍 Preview build  

17 files changed · + 4 added · ± 13 modified

+ Added

• api/dims/generated/pymc.dims.WeightedZeroSumNormal.html
• api/distributions/generated/pymc.WeightedZeroSumNormal.html
• api/dims/generated/classmethods/pymc.dims.WeightedZeroSumNormal.dist.html
• api/distributions/generated/classmethods/pymc.WeightedZeroSumNormal.dist.html

± Modified

• genindex.html
• glossary.html
• api/dims.html
• api/distributions.html
• api/dims/distributions.html
• api/distributions/multivariate.html
• _modules/pymc/distributions/multivariate.html
• _modules/pymc/distributions/transforms.html
• api/dims/generated/pymc.dims.ZeroSumNormal.html
• api/distributions/generated/pymc.WishartBartlett.html
• and 3 more...

[codecov]

Codecov Report

❌ Patch coverage is 94.73684% with 7 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.89%. Comparing base (da8fc47) to head (d394c9b).

Files with missing lines	Patch %	Lines
pymc/dims/distributions/vector.py	86.48%	5 Missing ⚠️
pymc/dims/distributions/transforms.py	96.77%	1 Missing ⚠️
pymc/distributions/transforms.py	95.45%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #8410      +/-   ##
==========================================
+ Coverage   91.87%   91.89%   +0.01%     
==========================================
  Files         128      128              
  Lines       21256    21386     +130     
==========================================
+ Hits        19530    19653     +123     
- Misses       1726     1733       +7     

Files with missing lines	Coverage Δ
pymc/distributions/__init__.py	100.00% <ø> (ø)
pymc/distributions/multivariate.py	95.59% <100.00%> (+0.20%)	⬆️
pymc/dims/distributions/transforms.py	99.27% <96.77%> (-0.73%)	⬇️
pymc/distributions/transforms.py	99.59% <95.45%> (-0.41%)	⬇️
pymc/dims/distributions/vector.py	88.69% <86.48%> (-1.31%)	⬇️

🚀 New features to boost your workflow:

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

[juanitorduz]

Thanks for the PR — the math is right and cleanly implemented. I checked the Householder construction, the isometry claim, exact parity with ZeroSumTransform for equal weights, and that the logp matches a reference MvNormal(0, σ²I_{n-1}) on the isometric coordinates (to ~1e-15). The new tests pass locally.

That said, I don't think it is mergeable as-is: there are three functional bugs in exactly the regimes the PR description advertises (batched sigma, and symbolic/Data weights for the MMM use case — including any model copy/clone/freeze with such weights), plus a class of silent shape mismatches. The tests are green because they never exercise those paths. Everything marked verified below I reproduced by running this branch.

Blocking

1. Batched sigma breaks the logp (both APIs); with batch size == n it is silently wrong

weighted_zerosumnormal_logp in pymc/distributions/multivariate.py:

def weighted_zerosumnormal_logp(op, values, rng, size, sigma, weights, **kwargs):
    ...
    out = pt.sum(-0.5 * pt.pow(value / sigma, 2) - (...) * (n - 1) / n, axis=-1)

rv_op pads sigma with pt.shape_padright(sigma) before drawing, but the logp receives the raw sigma input (inputs[2]), so value of shape (3, n) divided by sigma of shape (3,) aligns sigma against the constrained axis. ZeroSumNormal avoids this only because its dist() runs pt.atleast_Nd(sigma, n=n_zerosum_axes) first.

Verified:

• pm.WeightedZeroSumNormal.dist(sigma=[1, 2, 3], weights=w) → draws are fine (3, 4), but pm.logp(...) raises ValueError: Incompatible Elemwise input shapes [(3, 4), (1, 3)].
• sigma=[1, 2, 3, 4] (batch == n): no error, wrong numbers — logp = [-5.23, -5.39, -5.20, -5.28] vs reference [-3.75, -5.71, -6.10, -6.96].
• dims API: pmd.WeightedZeroSumNormal("x", sigma=as_xtensor([1, 2, 3], dims=("a",)), weights=w, core_dims="b") → same ValueError from compile_logp.

test_batched_sigma only checks draws, never the logp. Fix (verified): sigma = pt.shape_padright(sigma) at the top of the logp (or pad in dist like ZSN does). With that, m.logp(sum=False)[0].type.shape == (3,) and the values match the reference. Please also add the analogues of the existing TestZeroSumNormal.test_batched_transformed_logp_shape and the dims test_zerosumnormal_batch_sigma — either would have caught this.

2. pm.Data weights in the tensor API cannot be sampled

WeightedZeroSumTransform.forward in pymc/distributions/transforms.py:

coef = value[..., -1:] / (1 + u_last)
return value[..., :-1] - coef * u_head

With symbolic weights the RV (and its value variable) has static shape (None,), so value[..., -1:] also has static shape (None,), and PyTensor refuses to broadcast a runtime-1 dimension that isn't statically 1. backward happens to be fine (keepdims=True yields a static 1), so pm.logp works, but forward is used by initial_point, so model.initial_point() and pm.sample raise ValueError: Runtime broadcasting not allowed.

Verified with pm.Data("w", w). The dims version is fine because isel({dim: -1}) drops the dimension. The PR description only claims symbolic weights for the dims API, but the motivating use case (spend shares as data in an MMM, which lives in the tensor API in pymc-marketing) is exactly this. Fix (verified): coef = pt.expand_dims(value[..., -1], -1) / (1 + u_last) (or pt.specify_broadcastable). A tensor-API pm.Data + set_data test mirroring the dims one would cover it.

3. The transform captures the weights graph variable instead of reading it from rv_inputs

weighted_zerosum_default_transform passes rv.owner.inputs[3] into WeightedZeroSumTransform(weights), and forward/backward/_weight_direction then use self.weights and ignore the *rv_inputs they are handed (the Interval/args_fn convention reads parameters from rv_inputs precisely to avoid this). The transform is therefore bound to one specific graph. Verified consequences, in both APIs:

• RV-valued weights: w = pm.Dirichlet("w", a=np.ones(4)); pm.WeightedZeroSumNormal("x", weights=w) → compile_logp() raises ValueError: Random variables detected in the logp graph: {w}, because backward re-inserts the original RV rather than its value variable.
• pm.Data weights: clone_model(m), freeze_dims_and_data(m) and m.copy() all raise ValueError: x is not in list (the transform_values rewrite fails with a MissingInputError on w, since the cloned graph deep-copies the shared variable but reuses the transform object pointing at the old one). ZeroSumNormal in the same setup clones fine. This also affects pm.do/pm.observe, which go through the same fgraph machinery.

Fix: read the weights from rv_inputs[3] inside _weight_direction (the dims transform equivalently), keeping the constructor argument only as a fallback for standalone use. Please add a clone_model/freeze_dims_and_data test and an RV-valued-weights test.

Should fix before merge

4. len(weights) is never validated against shape / dims / coords

• Tensor API: pm.WeightedZeroSumNormal.dist(weights=w, shape=(5, 7)) → shape (5, 4); the 7 is silently ignored. dims=("a", "b") with len(b) = 5 and 4 weights → the RV is created with shape (2, 4) and the failure only surfaces later as RuntimeError: Could not create variable during InferenceData conversion in pm.sample.
• dims API: coords={"b": range(5)} with 4 weights → RV created, draw shape (4,). dist() already receives dim_lengths but doesn't use it.

ZeroSumNormal avoids this by deriving support_shape from shape/dims. Here the weights are the support shape, so dist should assert weights.shape[0] == shape[-1] (static check when both are constants, Assert otherwise), and the dims dist should compare against dim_lengths[core_dims[0]].

5. No positivity / degeneracy validation for weights in the tensor API

The dims transform rejects weights <= 0 for constants; the tensor dist, transform and logp check nothing. Verified: negative, zero and all-zero weights all "work" (all-zero gets constant-folded by PyTensor into the unconstrained leading components — no NaN, no error). weights=[0, 0, -1] gives u_n = -1, i.e. division by zero in 1 + u_n, which only surfaces as a misleading ParameterValueError: sum(weights * value) = 0. Since the docstring says "strictly positive", it should be enforced: raise in dist for constants and add pt.all(weights > 0) to the check_parameters call in the logp (which also covers the symbolic case). Unlike pm.Normal, the logp also has no sigma > 0 check (sigma=-1 → nan); ZSN shares that, so it's parity, but worth adding while you're there.

6. The on-constraint tolerance is absolute, not scale-aware — float32 is effectively unusable beyond σ ≈ 10

atol = 1e-9 if float64 else 1e-6 on sum(value * u). Verified with the PR's own transform: float32 with sigma=1 passes 100% of points; with sigma=100 only 33% pass (max |u·x| = 3.6e-5), so the logp is -inf for most of the posterior and NUTS cannot move. float64 at sigma=1e6 already drops 0.15% of points (ZSN's mean-based check survives there with max 2.3e-10). Since the PR text specifically advertises the dtype-aware tolerance, it should be relative, e.g. atol * (1 + pt.abs(value).sum(-1)) or pt.isclose(..., rtol=...) against pt.linalg.norm(value, axis=-1). test_floatX only evaluates at σ = 1.

7. Several tests are vacuous

• test_weighted_zerosumnormal_matches_zerosumnormal_with_equal_weights (dims) compares logp at the initial point, which is the zero vector — it only checks the normalizing constant, not the quadratic term or the transform.
• test_weighted_zerosumnormal_symbolic_weights checks isfinite(logp(initial_point)); zeros satisfy every constraint, so the set_data sync claimed in the comment is never tested. Evaluate at a random unconstrained point, call set_data, and check that backward(z) satisfies the new constraint (it does — the test just isn't there).
• No logp test against an independent reference: an (n-1)-dim MvNormal evaluated at forward(x) is a two-line strong test (TestZeroSumNormal.test_zsn_logp does the equivalent).
• Missing: batched-sigma logp, pm.Data in the tensor API, shape/dims length mismatch errors, non-positive weights, and the dims error branches (weights=None, dims mismatch — the uncovered lines in codecov).
• test_sampling_stays_on_constraint runs NUTS to test what test_default_transform_bijective already tests.
• Nothing calls the tensor WeightedZeroSumTransform.log_jac_det (only the dims test does), and there is no support-point test (test_zerosum_normal_support_point is the precedent).

Design / documentation

8. Please state explicitly which weighted generalization this is

There are two natural "weighted zero-sum normals" with the same constraint Σ wᵢxᵢ = 0 but different covariances: (a) the orthogonal projection σ²(I − uuᵀ) (this PR; isotropic on the hyperplane), and (b) "subtract the weighted mean", x = z − (wᵀz / Σw)·1 (oblique projection, non-isotropic). They differ materially: for w = [0.9, 0.05, 0.03, 0.02], (a) gives Var(x₁) ≈ 0.004σ² while (b) gives ≈ 0.014σ², and Var(x₂) ≈ σ² vs ≈ 1.7σ². The choice in this PR is the right one for a clean logp (note the density doesn't depend on the weights at all except through the constraint), but the docstring should state the marginal variances σ²(1 − uᵢ²) — otherwise users will be surprised that the dominant unit is pinned near zero.

9. API questions for maintainers

• Single, last-axis constraint only. Fine, but ZeroSumNormal accepts n_zerosum_axes, and users will eventually expect a weights kwarg there. Adding it later means a second code path to maintain, so the "fold into ZSN vs. new class" decision deserves an explicit call now (the size = sigma.shape inference and the () sigma core mean the RV signatures genuinely differ, so folding is not a small change).
• The 6.5–8.7× ESS/sec claim in the commit message has no accompanying script; either drop it or link a reproducible benchmark.
• docs/source/api/distributions/multivariate.rst lists WeightedZeroSumNormal after WishartBartlett; alphabetical (and consistent with __init__.py) is before Wishart.

Nits

• WeightedZeroSumTransform is not added to pymc/distributions/transforms.py::__all__ (ZeroSumTransform is), nor to docs/source/api/distributions/transforms.rst / docs/source/api/dims/transforms.rst.
• The logp is a copy of zerosumnormal_logp's single-axis density differing only in the constraint check, and weights / sqrt(sum(weights**2)) is hand-rolled at four sites (RV, logp, both transforms). A shared helper would keep the RV, logp and transform bit-identical on u, which the atol=1e-9 check depends on; the dims __new__ is likewise a copy of the ZeroSumNormal one (down to the dim_lenghts typo in the comment).
• The tensor WeightedZeroSumTransform.__init__ re-casts weights that dist already cast to floatX; harmless.
• test_default_transform_bijective imports inside the function.
• initvals={"x": ...} with the dims transform raises AttributeError: 'TensorConstant' object has no attribute 'isel' — this is pre-existing in the dims ZeroSumTransform too (.sizes), so not this PR's responsibility, but worth a separate issue.

Checked and correct

• Householder H = I − 2vvᵀ/vᵀv with v = u + eₙ maps u → −eₙ; forward/backward are exactly (Hx)[:-1] and H[z, 0]; the sign choice keeps 1 + uₙ ≥ 1 for positive weights.
• Isometry ⇒ log_jac_det = 0, matching ZSN's convention; equal weights reproduce ZeroSumTransform algebraically.
• Logp −‖x‖²/(2σ²) − (n−1) log(√(2π) σ) is the correct density on the hyperplane for scalar sigma.
• change_dist_size, observed, sample_prior_predictive, shape=/size= batching, dims set_data re-sync, and the dims error branches all behave correctly.

Items 1–2 are one-line fixes and 3 is a small refactor of the transforms; 4–7 are small but I'd want them in the same PR.