Add cat constraint and CatTransform

numpyro/distributions/constraints.py

@@ -28,6 +28,7 @@
 
 __all__ = [
     "boolean",
+    "cat",
     "circular",
     "complex",
     "corr_cholesky",
@@ -65,7 +66,7 @@
 ]
 
 import math
-from typing import ClassVar, Generic, Optional, cast
+from typing import ClassVar, Generic, Optional, Sequence, cast
 
 import numpy as np
 
@@ -427,6 +428,105 @@ def eq(self, other: object, static: bool = False) -> ArrayLike:
         return self.base_constraint.eq(other.base_constraint, static=static)
 
 
+class _Cat(Constraint[NonScalarArray]):
+    """
+    Applies a sequence of constraints to slices along a dimension, in a way
+    compatible with :func:`jax.numpy.concatenate`.
+
+    :param cseq: Sequence of constraints to apply to consecutive slices.
+    :param int dim: Dimension along which to slice.
+    :param lengths: Length of each slice. Defaults to one per constraint.
+    """
+
+    def __init__(
+        self,
+        cseq: Sequence[Constraint],
+        dim: int = 0,
+        lengths: Optional[Sequence[int]] = None,
+    ) -> None:
+        assert cseq, "cseq cannot be empty"
+        assert all(isinstance(c, Constraint) for c in cseq), (
+            "cseq must contain only Constraint instances"
+        )
+        assert isinstance(dim, int), "dim must be an integer"
+        self.cseq = tuple(cseq)
+        if lengths is None:
+            lengths = (1,) * len(self.cseq)
+        assert len(lengths) == len(self.cseq), (
+            "lengths must have the same number of elements as cseq"
+        )
+        assert all(isinstance(length, int) and length >= 0 for length in lengths), (
+            "lengths must contain only nonnegative integers"
+        )
+        self.lengths = tuple(lengths)
+        self.dim = dim
+
+    @property
+    def is_discrete(self) -> bool:
+        return any(c.is_discrete for c in self.cseq)
+
+    @property
+    def event_dim(self) -> int:
+        return max(c.event_dim for c in self.cseq)
+
+    def _slices(self, value: NonScalarArray) -> list[NonScalarArray]:
+        ndim = jnp.ndim(value)
+        if not -ndim <= self.dim < ndim:
+            raise ValueError(
+                f"dim {self.dim} out of range for value with {ndim} dimensions"
+            )
+        if jnp.shape(value)[self.dim] != sum(self.lengths):
+            raise ValueError(
+                f"value.shape[{self.dim}] = {jnp.shape(value)[self.dim]} must equal "
+                f"the sum of lengths {sum(self.lengths)}"
+            )
+
+        values = []
+        start = 0
+        for length in self.lengths:
+            values.append(
+                jax.lax.slice_in_dim(value, start, start + length, axis=self.dim)
+            )
+            start += length
+        return values
+
+    def __call__(self, x: NonScalarArray) -> ArrayLike:
+        checks = [
+            constraint(value) for constraint, value in zip(self.cseq, self._slices(x))
+        ]
+        return jnp.concatenate(checks, axis=self.dim)
+
+    def feasible_like(self, prototype: NonScalarArray) -> NonScalarArray:
+        values = [
+            constraint.feasible_like(value)
+            for constraint, value in zip(self.cseq, self._slices(prototype))
+        ]
+        return jnp.concatenate(values, axis=self.dim)
+
+    def __repr__(self) -> str:
+        return "{}({}, dim={}, lengths={})".format(
+            self.__class__.__name__[1:], self.cseq, self.dim, self.lengths
+        )
+
+    def tree_flatten(self):
+        return (self.cseq,), (
+            ("cseq",),
+            {"dim": self.dim, "lengths": self.lengths},
+        )
+
+    def eq(self, other: object, static: bool = False) -> ArrayLike:
+        if not isinstance(other, _Cat):
+            return False
+        if self.dim != other.dim or self.lengths != other.lengths:
+            return False
+        if static:
+            return all(c1.eq(c2, static=True) for c1, c2 in zip(self.cseq, other.cseq))
+        result = jnp.array(True)
+        for c1, c2 in zip(self.cseq, other.cseq):
+            result = result & c1.eq(c2, static=False)
+        return result
+
+
 class _RealVector(
     _IndependentConstraint[NonScalarArray], _SingletonConstraint[NonScalarArray]
 ):
@@ -874,6 +974,7 @@ def tree_flatten(self):
 
 
 boolean = _Boolean()
+cat = _Cat
 circular = _Circular()
 complex = _Complex()
 corr_cholesky = _CorrCholesky()

numpyro/distributions/transforms.py

@@ -41,6 +41,7 @@
     "biject_to",
     "AbsTransform",
     "AffineTransform",
+    "CatTransform",
     "CholeskyTransform",
     "ComplexTransform",
     "ComposeTransform",
@@ -454,6 +455,163 @@ def eq(self, other: object, static: bool = False) -> ArrayLike:
         return result
 
 
+class CatTransform(Transform[NonScalarArray]):
+    """
+    Applies a sequence of transforms to consecutive slices along a dimension,
+    in a way compatible with :func:`jax.numpy.concatenate`.
+
+    :param tseq: Sequence of transforms to apply to consecutive slices.
+    :param int dim: Dimension along which to slice.
+    :param lengths: Length of each slice. Defaults to one per transform.
+    """
+
+    def __init__(
+        self,
+        tseq: Sequence[Transform],
+        dim: int = 0,
+        lengths: Optional[Sequence[int]] = None,
+    ) -> None:
+        assert tseq, "tseq cannot be empty"
+        assert all(isinstance(t, Transform) for t in tseq), (
+            "tseq must contain only Transform instances"
+        )
+        assert isinstance(dim, int), "dim must be an integer"
+        self.transforms = tuple(tseq)
+        if lengths is None:
+            lengths = (1,) * len(self.transforms)
+        assert len(lengths) == len(self.transforms), (
+            "lengths must have the same number of elements as tseq"
+        )
+        assert all(isinstance(length, int) and length >= 0 for length in lengths), (
+            "lengths must contain only nonnegative integers"
+        )
+        self.lengths = tuple(lengths)
+        self.dim = dim
+
+    @property
+    def length(self) -> int:
+        return sum(self.lengths)
+
+    @property
+    def domain(self) -> Constraint:
+        return constraints.cat(
+            [transform.domain for transform in self.transforms],
+            self.dim,
+            self.lengths,
+        )
+
+    @property
+    def codomain(self) -> Constraint:
+        return constraints.cat(
+            [transform.codomain for transform in self.transforms],
+            self.dim,
+            self.lengths,
+        )
+
+    def _slices(self, value: NonScalarArray) -> list[NonScalarArray]:
+        ndim = jnp.ndim(value)
+        if not -ndim <= self.dim < ndim:
+            raise ValueError(
+                f"dim {self.dim} out of range for value with {ndim} dimensions"
+            )
+        if jnp.shape(value)[self.dim] != self.length:
+            raise ValueError(
+                f"value.shape[{self.dim}] = {jnp.shape(value)[self.dim]} must equal "
+                f"the sum of lengths {self.length}"
+            )
+
+        values = []
+        start = 0
+        for length in self.lengths:
+            values.append(lax.slice_in_dim(value, start, start + length, axis=self.dim))
+            start += length
+        return values
+
+    def __call__(self, x: NonScalarArray) -> NonScalarArray:
+        return jnp.concatenate(
+            [
+                transform(value)
+                for transform, value in zip(self.transforms, self._slices(x))
+            ],
+            axis=self.dim,
+        )
+
+    def _inverse(self, y: NonScalarArray) -> NonScalarArray:
+        return jnp.concatenate(
+            [
+                transform.inv(value)
+                for transform, value in zip(self.transforms, self._slices(y))
+            ],
+            axis=self.dim,
+        )
+
+    def log_abs_det_jacobian(
+        self,
+        x: NonScalarArray,
+        y: NonScalarArray,
+        intermediates: Optional[PyTree] = None,
+    ) -> NumLike:
+        if intermediates is not None and len(intermediates) != len(self.transforms):
+            raise ValueError(
+                f"Intermediates array has length = {len(intermediates)}. "
+                f"Expected = {len(self.transforms)}."
+            )
+
+        event_dim = max(self.domain.event_dim, self.codomain.event_dim)
+        logdetjacs = []
+        for i, (transform, xslice, yslice) in enumerate(
+            zip(self.transforms, self._slices(x), self._slices(y))
+        ):
+            intermediate = None if intermediates is None else intermediates[i]
+            logdetjac = transform.log_abs_det_jacobian(
+                xslice, yslice, intermediates=intermediate
+            )
+            transform_event_dim = max(
+                transform.domain.event_dim, transform.codomain.event_dim
+            )
+            logdetjacs.append(sum_rightmost(logdetjac, event_dim - transform_event_dim))
+
+        dim = self.dim
+        if dim >= 0:
+            dim -= jnp.ndim(x)
+        dim += event_dim
+        if dim < 0:
+            return jnp.concatenate(logdetjacs, axis=dim)
+        return sum(logdetjacs)
+
+    def call_with_intermediates(
+        self, x: NonScalarArray
+    ) -> Tuple[NumLike, Optional[PyTree]]:
+        values = []
+        intermediates = []
+        for transform, value in zip(self.transforms, self._slices(x)):
+            value, intermediate = transform.call_with_intermediates(value)
+            values.append(value)
+            intermediates.append(intermediate)
+        return jnp.concatenate(values, axis=self.dim), intermediates
+
+    def tree_flatten(self):
+        return (self.transforms,), (
+            ("transforms",),
+            {"dim": self.dim, "lengths": self.lengths},
+        )
+
+    def eq(self, other: object, static: bool = False) -> ArrayLike:
+        if not isinstance(other, CatTransform):
+            return False
+        if self.dim != other.dim or self.lengths != other.lengths:
+            return False
+        if static:
+            return all(
+                t1.eq(t2, static=True)
+                for t1, t2 in zip(self.transforms, other.transforms)
+            )
+        result = jnp.array(True)
+        for t1, t2 in zip(self.transforms, other.transforms):
+            result = result & t1.eq(t2, static=False)
+        return result
+
+
 def _matrix_forward_shape(shape: tuple[int, ...], offset: int = 0) -> tuple[int, ...]:
     # Reshape from (..., N) to (..., D, D).
     if len(shape) < 1:
@@ -2100,6 +2258,13 @@ def __call__(self, constraint):
 biject_to = ConstraintRegistry()
 
 
+@biject_to.register(constraints.cat)
+def _biject_to_cat(constraint):
+    return CatTransform(
+        [biject_to(c) for c in constraint.cseq], constraint.dim, constraint.lengths
+    )
+
+
 @biject_to.register(constraints.corr_cholesky)
 def _transform_to_corr_cholesky(constraint):
     return CorrCholeskyTransform()

test/test_distributions.py

@@ -2649,6 +2649,38 @@ def test_beta_proportion_invalid_mean():
         (constraints.boolean, np.array([True, False]), np.array([True, True])),
         (constraints.boolean, np.array([1, 1]), np.array([True, True])),
         (constraints.boolean, np.array([-1, 1]), np.array([False, True])),
+        (
+            constraints.cat(
+                [constraints.interval(-1, 1), constraints.positive],
+                dim=-1,
+                lengths=[2, 1],
+            ),
+            np.array([[0.0, 2.0, 1.0], [-2.0, 0.5, -1.0]]),
+            np.array([[True, False, True], [False, True, False]]),
+        ),
+        (
+            constraints.cat([constraints.positive, constraints.unit_interval]),
+            np.array([[1.0, 0.0, -1.0], [0.0, 0.5, 2.0]]),
+            np.array([[True, False, False], [True, True, False]]),
+        ),
+        (
+            constraints.cat(
+                [constraints.less_than(0), constraints.nonnegative],
+                dim=1,
+                lengths=[1, 2],
+            ),
+            np.array([[-1.0, 0.0, 2.0], [1.0, -1.0, 0.0]]),
+            np.array([[True, True, True], [False, False, True]]),
+        ),
+        (
+            constraints.cat(
+                [constraints.positive, constraints.unit_interval],
+                dim=-1,
+                lengths=[0, 2],
+            ),
+            np.array([[0.0, 0.5], [-1.0, 2.0]]),
+            np.array([[True, True], [False, False]]),
+        ),
         (
             constraints.corr_cholesky,
             np.array([[[1, 0], [0, 1]], [[1, 0.1], [0, 1]]]),
@@ -2783,6 +2815,30 @@ def test_constraints(constraint, x, expected):
         assert_allclose(inverse, jnp.zeros_like(inverse), atol=2e-7)
 
 
+def test_cat_constraint_pytree_and_validation():
+    constraint = constraints.cat(
+        [constraints.interval(-1.0, 1.0), constraints.positive],
+        dim=-1,
+        lengths=[2, 1],
+    )
+    value = jnp.array([[0.0, 2.0, 1.0], [-2.0, 0.5, -1.0]])
+    expected = jnp.array([[True, False, True], [False, True, False]])
+
+    assert_array_equal(jax.jit(lambda c, x: c(x))(constraint, value), expected)
+    leaves, treedef = jax.tree.flatten(constraint)
+    assert constraint.eq(jax.tree.unflatten(treedef, leaves), static=True)
+
+    with pytest.raises(ValueError, match="must equal the sum of lengths 3"):
+        constraint(jnp.ones(2))
+
+    with pytest.raises(AssertionError, match="cseq cannot be empty"):
+        constraints.cat([])
+    with pytest.raises(AssertionError, match="dim must be an integer"):
+        constraints.cat([constraints.real], dim=0.5)
+    with pytest.raises(AssertionError, match="nonnegative integers"):
+        constraints.cat([constraints.real], lengths=[-1])
+
+
 @pytest.mark.parametrize(
     "constraint",
     [