From b46eefba1cb4104f6100fc91c3ad5e9ab903b1f2 Mon Sep 17 00:00:00 2001 From: ricardoV94 Date: Fri, 27 Feb 2026 20:23:12 +0100 Subject: [PATCH 1/4] Implement dims Censored --- .github/workflows/tests.yml | 1 + docs/source/api/dims/distributions.rst | 11 +++ pymc/dims/distributions/__init__.py | 1 + pymc/dims/distributions/censored.py | 51 +++++++++++++ pymc/dims/distributions/core.py | 26 ++++++- pymc/dims/distributions/vector.py | 4 +- tests/dims/distributions/test_censored.py | 88 +++++++++++++++++++++++ 7 files changed, 179 insertions(+), 3 deletions(-) create mode 100644 pymc/dims/distributions/censored.py create mode 100644 tests/dims/distributions/test_censored.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index df5ff8d10b..655776e963 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -139,6 +139,7 @@ jobs: - | tests/dims/distributions/test_core.py + tests/dims/distributions/test_censored.py tests/dims/distributions/test_scalar.py tests/dims/distributions/test_vector.py tests/dims/test_model.py diff --git a/docs/source/api/dims/distributions.rst b/docs/source/api/dims/distributions.rst index 229980b73a..cc34d68deb 100644 --- a/docs/source/api/dims/distributions.rst +++ b/docs/source/api/dims/distributions.rst @@ -39,3 +39,14 @@ Vector distributions Categorical MvNormal ZeroSumNormal + + +Higher-Order distributions +========================== + +.. currentmodule:: pymc.dims +.. autosummary:: + :toctree: generated/ + :template: distribution.rst + + Censored diff --git a/pymc/dims/distributions/__init__.py b/pymc/dims/distributions/__init__.py index da85bc1463..6c49789089 100644 --- a/pymc/dims/distributions/__init__.py +++ b/pymc/dims/distributions/__init__.py @@ -11,5 +11,6 @@ # 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. +from pymc.dims.distributions.censored import Censored from pymc.dims.distributions.scalar import * from pymc.dims.distributions.vector import * diff --git a/pymc/dims/distributions/censored.py b/pymc/dims/distributions/censored.py new file mode 100644 index 0000000000..29d9a56136 --- /dev/null +++ b/pymc/dims/distributions/censored.py @@ -0,0 +1,51 @@ +# 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. +import numpy as np + +from pymc.dims.distributions.core import DimDistribution, copy_docstring, expand_dist_dims +from pymc.distributions.censored import Censored as RegularCensored + + +@copy_docstring(RegularCensored) +class Censored(DimDistribution): + @classmethod + def dist(cls, dist, *, lower=None, upper=None, dim_lengths, **kwargs): + if lower is None: + lower = -np.inf + if upper is None: + upper = np.inf + return super().dist([dist, lower, upper], dim_lengths=dim_lengths, **kwargs) + + @classmethod + def xrv_op(cls, dist, lower, upper, core_dims=None, extra_dims=None, rng=None): + if extra_dims is None: + extra_dims = {} + + dist = cls._as_xtensor(dist) + lower = cls._as_xtensor(lower) + upper = cls._as_xtensor(upper) + + # Any dimensions in extra_dims, or only present in lower, upper, + # must propagate back to the dist as `extra_dims` + bounds_sizes = lower.sizes | upper.sizes + dist_dims_set = set(dist.dims) + extra_dist_dims = extra_dims | { + dim: size for dim, size in bounds_sizes.items() if dim not in dist_dims_set + } + if extra_dist_dims: + dist = expand_dist_dims(dist, extra_dist_dims) + + # Probability is inferred from the clip operation + # TODO: Make this a SymbolicRandomVariable that can itself be resized + return dist.clip(lower, upper) diff --git a/pymc/dims/distributions/core.py b/pymc/dims/distributions/core.py index aee8e4cf7f..918b28a794 100644 --- a/pymc/dims/distributions/core.py +++ b/pymc/dims/distributions/core.py @@ -13,7 +13,7 @@ # limitations under the License. from collections.abc import Callable, Sequence from itertools import chain -from typing import cast +from typing import Any, cast import numpy as np @@ -25,7 +25,9 @@ from pytensor.tensor.random.op import RandomVariable from pytensor.xtensor import as_xtensor from pytensor.xtensor.basic import XTensorFromTensor, xtensor_from_tensor +from pytensor.xtensor.shape import Transpose from pytensor.xtensor.type import XTensorVariable +from pytensor.xtensor.vectorization import XRV from pymc import SymbolicRandomVariable, modelcontext from pymc.dims.distributions.transforms import DimTransform, log_odds_transform, log_transform @@ -345,3 +347,25 @@ class UnitDimDistribution(DimDistribution): """Base class for unit-valued distributions.""" default_transform = log_odds_transform + + +def expand_dist_dims(dist: XTensorVariable, extra_dims: dict[str, Any]) -> XTensorVariable: + if overlap := (set(extra_dims) & set(dist.dims)): + raise ValueError(f"extra_dims already present in distribution: {sorted(overlap)}") + + op = None if dist.owner is None else dist.owner.op + match op: + case XRV(): + # Recreate dist with new extra dims + dist_props = dist.owner.op._props_dict() + dist_props["extra_dims"] = (*(extra_dims.keys()), *dist_props["extra_dims"]) + new_dist_op = type(dist.owner.op)(**dist_props) + _old_rng, *params_and_dim_lengths = dist.owner.inputs + new_rng = None # We don't propagate the old RNG, because we don't want the new and old dists to be correlated + return new_dist_op(new_rng, *extra_dims.values(), *params_and_dim_lengths) + case Transpose(): + return expand_dist_dims(dist.owner.inputs[0], extra_dims=extra_dims).transpose( + ..., *dist.dims + ) + case _: + raise NotImplementedError(f"expand_dist_dims not implemented for {dist} with op {op}") diff --git a/pymc/dims/distributions/vector.py b/pymc/dims/distributions/vector.py index 7107712e67..0875afb6e0 100644 --- a/pymc/dims/distributions/vector.py +++ b/pymc/dims/distributions/vector.py @@ -238,8 +238,8 @@ def dist(cls, sigma=1.0, *, core_dims=None, dim_lengths, **kwargs): ) @classmethod - def xrv_op(self, sigma, support_dims, core_dims, extra_dims=None, rng=None): - sigma = as_xtensor(sigma) + def xrv_op(cls, sigma, support_dims, core_dims, extra_dims=None, rng=None): + sigma = cls._as_xtensor(sigma) support_dims = as_xtensor(support_dims, dims=("_",)) support_shape = support_dims.values core_rv = ZeroSumNormalRV.rv_op(sigma=sigma.values, support_shape=support_shape).owner.op diff --git a/tests/dims/distributions/test_censored.py b/tests/dims/distributions/test_censored.py new file mode 100644 index 0000000000..eb96f7a4da --- /dev/null +++ b/tests/dims/distributions/test_censored.py @@ -0,0 +1,88 @@ +# 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. +import numpy as np +import pytest + +from pytensor.xtensor import as_xtensor +from pytensor.xtensor.shape import Transpose +from pytensor.xtensor.vectorization import XRV + +import pymc.distributions as regular_distributions + +from pymc.dims import Censored, Normal +from pymc.model.core import Model +from tests.dims.utils import assert_equivalent_logp_graph, assert_equivalent_random_graph + + +@pytest.mark.parametrize("lower", [None, -1]) +@pytest.mark.parametrize("upper", [None, 1]) +def test_censored_basic(lower, upper): + coords = {"space": range(3), "time": range(4)} + + with Model(coords=coords) as model: + dist = Normal.dist(np.pi, np.e) + Censored("y", dist, lower=lower, upper=upper, dims=("space", "time")) + + with Model(coords=coords) as reference_model: + dist = regular_distributions.Normal.dist(np.pi, np.e) + regular_distributions.Censored( + "y", dist=dist, lower=lower, upper=upper, dims=("space", "time") + ) + + assert_equivalent_random_graph(model, reference_model) + assert_equivalent_logp_graph(model, reference_model) + + +def test_censored_dims(): + """Test that both censored (and the underlying dist) have all the implied and explicit dims.""" + coords = { + "a": range(3), + "b": range(4), + "c": range(5), + "d": range(6), + } + with Model(coords=coords) as model: + dist = Normal.dist( + mu=as_xtensor([0, 1, 2], dims=("a",)), + sigma=as_xtensor([1, 2, 3], dims=("b",)), + dim_lengths={"c": model.dim_lengths["c"]}, + ) + assert set(dist.dims) == {"c", "a", "b"} + + c0 = Censored("c0", dist) + assert c0.dims == ("c", "a", "b") + c0_dist = c0.owner.inputs[0] + assert isinstance(c0_dist.owner.op, XRV) + assert c0_dist.dims == ("c", "a", "b") + + c1 = Censored("c1", dist, dims=("a", "b", "c")) + assert c1.dims == ("a", "b", "c") + assert isinstance(c1.owner.op, Transpose) + c1_dist = c1.owner.inputs[0].owner.inputs[0] + assert isinstance(c1_dist.owner.op, XRV) + assert c1_dist.dims == ("c", "a", "b") + + c2 = Censored("c2", dist, dims=(..., "d")) + assert c2.dims == ("c", "a", "b", "d") + assert isinstance(c1.owner.op, Transpose) + c2_dist = c2.owner.inputs[0].owner.inputs[0] + assert isinstance(c2_dist.owner.op, XRV) + assert c2_dist.dims == ("d", "c", "a", "b") + + lower = as_xtensor(np.zeros((6, 5)), dims=("d", "c")) + c3 = Censored("c3", dist, lower=lower) + assert c3.dims == ("d", "c", "a", "b") + c3_dist = c3.owner.inputs[0] + assert isinstance(c3_dist.owner.op, XRV) + assert c3_dist.dims == ("d", "c", "a", "b") From de84781dd14b66ec8181399eddd2377e8746a3bf Mon Sep 17 00:00:00 2001 From: ricardoV94 Date: Sat, 28 Feb 2026 16:19:51 +0100 Subject: [PATCH 2/4] ZeroSumNormal: Explicit sigma signature --- pymc/distributions/multivariate.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/pymc/distributions/multivariate.py b/pymc/distributions/multivariate.py index 192fff6e30..51d6f7f243 100644 --- a/pymc/distributions/multivariate.py +++ b/pymc/distributions/multivariate.py @@ -2677,8 +2677,10 @@ def rv_op(cls, sigma, support_shape, *, size=None, rng=None): for axis in range(n_zerosum_axes): zerosum_rv -= zerosum_rv.mean(axis=-axis - 1, keepdims=True) + # sigma has core_shape = (1, 1, ...) (as many as there are zerosum axes) + ones = ",".join("1" for _ in range(n_zerosum_axes)) support_str = ",".join([f"d{i}" for i in range(n_zerosum_axes)]) - extended_signature = f"[rng],[size],(),(s)->[rng],({support_str})" + extended_signature = f"[rng],[size],({ones}),(s)->[rng],({support_str})" return cls( inputs=[rng, size, sigma, support_shape], outputs=[next_rng, zerosum_rv], @@ -2774,7 +2776,7 @@ def __new__(cls, *args, n_zerosum_axes=None, support_shape=None, dims=None, **kw def dist(cls, sigma=1.0, n_zerosum_axes=None, support_shape=None, **kwargs): n_zerosum_axes = cls.check_zerosum_axes(n_zerosum_axes) - sigma = pt.as_tensor(sigma) + sigma = pt.atleast_Nd(pt.as_tensor(sigma), n=n_zerosum_axes) if not all(sigma.type.broadcastable[-n_zerosum_axes:]): raise ValueError("sigma must have length one across the zero-sum axes") From 2dd5c044ff6c944227389bfa1357ad890378abb6 Mon Sep 17 00:00:00 2001 From: ricardoV94 Date: Sat, 28 Feb 2026 04:34:25 +0100 Subject: [PATCH 3/4] ZeroSumNormal: Fix logp shape of transformed variates --- pymc/distributions/multivariate.py | 2 +- pymc/distributions/transforms.py | 2 +- tests/distributions/test_multivariate.py | 7 +++++++ 3 files changed, 9 insertions(+), 2 deletions(-) diff --git a/pymc/distributions/multivariate.py b/pymc/distributions/multivariate.py index 51d6f7f243..6f8a2abe26 100644 --- a/pymc/distributions/multivariate.py +++ b/pymc/distributions/multivariate.py @@ -2659,9 +2659,9 @@ class ZeroSumNormalRV(SymbolicRandomVariable): @classmethod def rv_op(cls, sigma, support_shape, *, size=None, rng=None): + support_shape = _squeeze_to_ndim(pt.as_tensor(support_shape), ndim=1) n_zerosum_axes = pt.get_vector_length(support_shape) sigma = pt.as_tensor(sigma) - support_shape = pt.as_tensor(support_shape, ndim=1) rng = normalize_rng_param(rng) size = normalize_size_param(size) diff --git a/pymc/distributions/transforms.py b/pymc/distributions/transforms.py index 5e3f934953..c56cba8a55 100644 --- a/pymc/distributions/transforms.py +++ b/pymc/distributions/transforms.py @@ -309,7 +309,7 @@ def backward(self, value, *rv_inputs): return value def log_jac_det(self, value, *rv_inputs): - return pt.constant(0.0) + return value.sum(self.zerosum_axes).zeros_like() log_exp_m1 = LogExpM1() diff --git a/tests/distributions/test_multivariate.py b/tests/distributions/test_multivariate.py index 67ac8644d0..27d145aecb 100644 --- a/tests/distributions/test_multivariate.py +++ b/tests/distributions/test_multivariate.py @@ -1763,6 +1763,13 @@ def test_batched_sigma(self): sigma=batch_test_sigma[None, :, None], n_zerosum_axes=2, support_shape=(3, 2) ) + def test_batched_transformed_logp_shape(self): + with pm.Model() as m: + x = pm.ZeroSumNormal("x", sigma=np.ones(3)[:, None], support_shape=(2,)) + assert x.type.shape == (3, 2) + assert m.logp(sum=False)[0].type.shape == (3,) + assert m.logp(sum=False, jacobian=False)[0].type.shape == (3,) + class TestMvStudentTCov(BaseTestDistributionRandom): def mvstudentt_rng_fn(self, size, nu, mu, scale, rng): From 03965c1021c1e0d59003d03162ef3ede99d94e1f Mon Sep 17 00:00:00 2001 From: ricardoV94 Date: Sat, 28 Feb 2026 04:41:44 +0100 Subject: [PATCH 4/4] ZeroSumNormal: Fix dims variant with batch sigma --- pymc/dims/distributions/transforms.py | 4 +--- pymc/dims/distributions/vector.py | 28 +++++++++++++++++-------- pymc/distributions/distribution.py | 2 +- tests/dims/distributions/test_vector.py | 20 ++++++++++++++++++ 4 files changed, 41 insertions(+), 13 deletions(-) diff --git a/pymc/dims/distributions/transforms.py b/pymc/dims/distributions/transforms.py index 14dbb9444f..e898ab0508 100644 --- a/pymc/dims/distributions/transforms.py +++ b/pymc/dims/distributions/transforms.py @@ -203,6 +203,4 @@ def backward(self, value, *rv_inputs): return value def log_jac_det(self, value, *rv_inputs): - # Use following once broadcast_like is implemented - # as_xtensor(0).broadcast_like(value, exclude=self.dims)` - return value.sum(self.dims) * 0 + return as_xtensor(0.0).broadcast_like(value, exclude=self.dims) diff --git a/pymc/dims/distributions/vector.py b/pymc/dims/distributions/vector.py index 0875afb6e0..0995f891e8 100644 --- a/pymc/dims/distributions/vector.py +++ b/pymc/dims/distributions/vector.py @@ -229,7 +229,8 @@ def dist(cls, sigma=1.0, *, core_dims=None, dim_lengths, **kwargs): raise ValueError("ZeroSumNormal requires atleast 1 core_dims") support_dims = as_xtensor( - as_tensor([dim_lengths[core_dim] for core_dim in core_dims]), dims=("_",) + as_tensor([dim_lengths[core_dim] for core_dim in core_dims]), + dims=("__support_shape__",), ) sigma = cls._as_xtensor(sigma) @@ -238,16 +239,25 @@ def dist(cls, sigma=1.0, *, core_dims=None, dim_lengths, **kwargs): ) @classmethod - def xrv_op(cls, sigma, support_dims, core_dims, extra_dims=None, rng=None): - sigma = cls._as_xtensor(sigma) - support_dims = as_xtensor(support_dims, dims=("_",)) - support_shape = support_dims.values - core_rv = ZeroSumNormalRV.rv_op(sigma=sigma.values, support_shape=support_shape).owner.op + def xrv_op(cls, sigma, support_shape, core_dims, extra_dims=None, rng=None): + # ZeroSumNormal expects dummy dimensions on sigma for the support_shape + sigma = cls._as_xtensor(sigma).expand_dims(core_dims) + support_shape = as_xtensor(support_shape, dims=("__support_shape__",)) + core_rv = ZeroSumNormalRV.rv_op( + sigma=sigma.values, support_shape=support_shape.values + ).owner.op + core_dims_map = tuple(range(1, len(core_dims) + 1)) xop = pxr.as_xrv( core_rv, - core_inps_dims_map=[(), (0,)], - core_out_dims_map=tuple(range(1, len(core_dims) + 1)), + core_inps_dims_map=[core_dims_map, (0,)], + core_out_dims_map=core_dims_map, ) # Dummy "_" core dim to absorb the support_shape vector # If ZeroSumNormal expected a scalar per support dim, this wouldn't be needed - return xop(sigma, support_dims, core_dims=("_", *core_dims), extra_dims=extra_dims, rng=rng) + return xop( + sigma, + support_shape, + core_dims=("__support_shape__", *core_dims), + extra_dims=extra_dims, + rng=rng, + ) diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py index b434a6db25..7fcbd75c83 100644 --- a/pymc/distributions/distribution.py +++ b/pymc/distributions/distribution.py @@ -392,7 +392,7 @@ def make_node(self, *inputs): ) if size_arg_idx is not None and len(rng_arg_idxs) == 1: new_size_type = normalize_size_param(inputs[size_arg_idx]).type - if not self.input_types[size_arg_idx].in_same_class(new_size_type): + if not self.input_types[size_arg_idx].is_super(new_size_type): params = [inputs[idx] for idx in param_idxs] size = inputs[size_arg_idx] rng = inputs[rng_arg_idxs[0]] diff --git a/tests/dims/distributions/test_vector.py b/tests/dims/distributions/test_vector.py index 8cfdadb372..a71ca2306d 100644 --- a/tests/dims/distributions/test_vector.py +++ b/tests/dims/distributions/test_vector.py @@ -99,3 +99,23 @@ def test_zerosumnormal(): # Logp is correct, but we have join(..., -1) and join(..., 1), that don't get canonicalized to the same # Should work once https://github.com/pymc-devs/pytensor/issues/1505 is fixed # assert_equivalent_logp_graph(model, reference_model) + + +def test_zerosumnormal_batch_sigma(): + coords = {"a": range(3), "b": range(5)} + sigma = np.array([1, 2, 3.0]) + with Model(coords=coords) as model: + ZeroSumNormal( + "x", + sigma=as_xtensor(sigma, dims=("a",)), + core_dims=("b",), + ) + + with Model(coords=coords) as ref_model: + regular_distributions.ZeroSumNormal("x", sigma=sigma[:, None], dims=("a", "b")) + + ip = model.initial_point() + np.testing.assert_allclose( + model.compile_logp(sum=False)(ip), + ref_model.compile_logp(sum=False)(ip), + )