diff --git a/numpyro/distributions/constraints.py b/numpyro/distributions/constraints.py index b71c46863..a0958e7cb 100644 --- a/numpyro/distributions/constraints.py +++ b/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() diff --git a/numpyro/distributions/transforms.py b/numpyro/distributions/transforms.py index 8a2c9b603..dc56ebd64 100644 --- a/numpyro/distributions/transforms.py +++ b/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() diff --git a/test/test_distributions.py b/test/test_distributions.py index fc9dade20..9e5cec28e 100644 --- a/test/test_distributions.py +++ b/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", [ diff --git a/test/test_transforms.py b/test/test_transforms.py index aaf540cb3..2b5a1994f 100644 --- a/test/test_transforms.py +++ b/test/test_transforms.py @@ -20,6 +20,7 @@ from numpyro.distributions.transforms import ( AbsTransform, AffineTransform, + CatTransform, CholeskyTransform, ComplexTransform, ComposeTransform, @@ -77,6 +78,11 @@ class T(namedtuple("TestCase", ["transform_cls", "params", "kwargs"])): ), dict(), ), + "cat": T( + CatTransform, + ([AffineTransform(np.array(1.0), np.array(2.0)), ExpTransform()],), + dict(dim=-1, lengths=(2, 1)), + ), "independent": T( IndependentTransform, (AffineTransform(np.array([1.0, 2.0]), np.array([3.0, 4.0])),), @@ -261,6 +267,55 @@ def test_reshape_transform_invalid(): ReshapeTransform((2, 3), (6,))(jnp.arange(2)) +def test_cat_transform(): + transform = CatTransform( + [AffineTransform(1.0, 2.0), ExpTransform()], dim=-1, lengths=[2, 1] + ) + x = jnp.array([[0.0, 1.0, 2.0], [-1.0, 3.0, 0.5]]) + expected = jnp.concatenate([1.0 + 2.0 * x[..., :2], jnp.exp(x[..., 2:])], -1) + + y, intermediates = jax.jit(transform.call_with_intermediates)(x) + assert jnp.allclose(y, expected) + assert jnp.allclose(jax.jit(lambda value: transform.inv(value))(y), x) + assert jnp.allclose( + jax.jit(transform.log_abs_det_jacobian)(x, y), + jnp.concatenate([jnp.full_like(x[..., :2], jnp.log(2.0)), x[..., 2:]], axis=-1), + ) + assert jnp.allclose( + transform.log_abs_det_jacobian(x, y, intermediates), + transform.log_abs_det_jacobian(x, y), + ) + + +def test_biject_to_cat_constraint(): + constraint = constraints.cat( + [constraints.interval(-2.0, 3.0), constraints.positive], + dim=-1, + lengths=[2, 1], + ) + transform = biject_to(constraint) + x = jnp.array([[0.0, -1.0, 2.0], [1.0, 0.5, -2.0]]) + y = jax.jit(lambda value: transform(value))(x) + + assert transform.codomain.dim == constraint.dim + assert transform.codomain.lengths == constraint.lengths + assert jnp.array_equal(constraint(y), jnp.ones_like(y, dtype=bool)) + assert jnp.allclose(jax.jit(lambda value: transform.inv(value))(y), x) + + +def test_cat_transform_invalid_shape(): + transform = CatTransform([ExpTransform(), ExpTransform()], lengths=[1, 2]) + with pytest.raises(ValueError, match="must equal the sum of lengths 3"): + transform(jnp.ones(2)) + + with pytest.raises(AssertionError, match="tseq cannot be empty"): + CatTransform([]) + with pytest.raises(AssertionError, match="dim must be an integer"): + CatTransform([ExpTransform()], dim=0.5) + with pytest.raises(AssertionError, match="nonnegative integers"): + CatTransform([ExpTransform()], lengths=[-1]) + + @pytest.mark.parametrize( "input_shape, shape, ndims", [