Merge main into v6

.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

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

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 *

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)

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}")

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)

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(self, sigma, support_dims, core_dims, extra_dims=None, rng=None):
-        sigma = 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,
+        )

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]]

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)
 
@@ -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")
 

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()

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")

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),
+    )

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):