Add Indexed, Permute and Reshape

distreqx/bijectors/__init__.py

@@ -7,7 +7,10 @@
 from ._block import Block as Block
 from ._chain import Chain as Chain
 from ._diag_linear import DiagLinear as DiagLinear
+from ._indexed import Indexed as Indexed
 from ._linear import AbstractLinearBijector as AbstractLinearBijector
+from ._permute import Permute as Permute
+from ._reshape import Reshape as Reshape
 from ._scalar_affine import ScalarAffine as ScalarAffine
 from ._shift import Shift as Shift
 from ._sigmoid import Sigmoid as Sigmoid

distreqx/bijectors/_indexed.py

@@ -0,0 +1,62 @@
+import equinox as eqx
+import jax.numpy as jnp
+from jaxtyping import Array
+
+from ._bijector import (
+    AbstractBijector,
+    AbstractForwardInverseBijector,
+    AbstractFwdLogDetJacBijector,
+    AbstractInvLogDetJacBijector,
+)
+
+
+class Indexed(
+    AbstractForwardInverseBijector,
+    AbstractInvLogDetJacBijector,
+    AbstractFwdLogDetJacBijector,
+):
+    """Applies a bijector to a specific subset of indices of an input array."""
+
+    bijector: AbstractBijector
+    indices: Array
+
+    _is_constant_jacobian: bool = eqx.field(init=False)
+    _is_constant_log_det: bool = eqx.field(init=False)
+
+    def __init__(self, bijector: AbstractBijector, indices: Array | list | tuple):
+        """Initializes an Indexed bijector.
+
+        **Arguments:**
+
+        - `bijector`: The bijector to apply to the specified subset.
+        - `indices`: An array of integer indices or boolean masks indicating
+            which elements of the input should be transformed.
+        """
+        self.bijector = bijector
+        self.indices = jnp.asarray(indices)
+
+        self._is_constant_jacobian = self.bijector.is_constant_jacobian
+        self._is_constant_log_det = self.bijector.is_constant_log_det
+
+    def forward_and_log_det(self, x: Array) -> tuple[Array, Array]:
+        """Transforms x at the target indices and leaves the rest unchanged."""
+        y_subset, log_det = self.bijector.forward_and_log_det(x[self.indices])
+        y = x.at[self.indices].set(y_subset)
+        return y, log_det
+
+    def inverse_and_log_det(self, y: Array) -> tuple[Array, Array]:
+        """
+        Inversely transforms y at the target indices and leaves the rest
+        unchanged.
+        """
+        x_subset, log_det = self.bijector.inverse_and_log_det(y[self.indices])
+        x = y.at[self.indices].set(x_subset)
+        return x, log_det
+
+    def same_as(self, other: AbstractBijector) -> bool:
+        """Returns True if this bijector is guaranteed to be the same as `other`."""
+        return bool(
+            type(other) is Indexed
+            and self.bijector.same_as(other.bijector)
+            and jnp.array_equal(self.indices, other.indices)
+        )

distreqx/bijectors/_permute.py

@@ -0,0 +1,65 @@
+import jax.numpy as jnp
+import numpy as np
+from jaxtyping import Array
+
+from ._bijector import (
+    AbstractBijector,
+    AbstractForwardInverseBijector,
+    AbstractFwdLogDetJacBijector,
+    AbstractInvLogDetJacBijector,
+)
+
+
+class Permute(
+    AbstractForwardInverseBijector,
+    AbstractInvLogDetJacBijector,
+    AbstractFwdLogDetJacBijector,
+):
+    """Permutation bijector that reorders the elements of a 1D event."""
+
+    permutation: Array
+    inverse_permutation: Array
+
+    _is_constant_jacobian: bool = True
+    _is_constant_log_det: bool = True
+
+    def __init__(self, permutation: Array | list | tuple):
+        """Initializes a Permute bijector.
+
+        **Arguments:**
+
+        - `permutation`: An array or list of integers representing the new order
+            of the elements. Must contain all integers from 0 to N-1 exactly once.
+        """
+        # Convert to numpy first to validate shapes/values before JIT compilation
+        perm_np = np.asarray(permutation, dtype=int)
+
+        if perm_np.ndim != 1:
+            raise ValueError(
+                f"Permutation must be a 1D array, got shape {perm_np.shape}."
+            )
+
+        expected_elements = np.arange(perm_np.size)
+        if not np.array_equal(np.sort(perm_np), expected_elements):
+            raise ValueError(
+                "Invalid permutation. It must contain all integers from "
+                f"0 to {perm_np.size - 1} exactly once."
+            )
+
+        self.permutation = jnp.asarray(perm_np)
+        self.inverse_permutation = jnp.argsort(self.permutation)
+
+    def forward_and_log_det(self, x: Array) -> tuple[Array, Array]:
+        """Computes y = x[permutation] and log|det J(f)(x)| = 0.0."""
+        return x[self.permutation], jnp.zeros((), dtype=x.dtype)
+
+    def inverse_and_log_det(self, y: Array) -> tuple[Array, Array]:
+        """Computes x = y[inverse_permutation] and log|det J(f^{-1})(y)| = 0.0."""
+        return y[self.inverse_permutation], jnp.zeros((), dtype=y.dtype)
+
+    def same_as(self, other: AbstractBijector) -> bool:
+        """Returns True if this bijector is guaranteed to be the same as `other`."""
+        return bool(
+            type(other) is Permute
+            and jnp.array_equal(self.permutation, other.permutation)
+        )

distreqx/bijectors/_reshape.py

@@ -0,0 +1,65 @@
+import math
+
+import equinox as eqx
+import jax.numpy as jnp
+from jaxtyping import Array
+
+from ._bijector import (
+    AbstractBijector,
+    AbstractForwardInverseBijector,
+    AbstractFwdLogDetJacBijector,
+    AbstractInvLogDetJacBijector,
+)
+
+
+class Reshape(
+    AbstractForwardInverseBijector,
+    AbstractInvLogDetJacBijector,
+    AbstractFwdLogDetJacBijector,
+):
+    """A bijector that reshapes the input array."""
+
+    in_shape: tuple[int, ...] = eqx.field(static=True)
+    out_shape: tuple[int, ...] = eqx.field(static=True)
+
+    _is_constant_jacobian: bool = True
+    _is_constant_log_det: bool = True
+
+    def __init__(self, in_shape: tuple[int, ...], out_shape: tuple[int, ...]):
+        """Initializes a Reshape bijector.
+
+        **Arguments:**
+
+        - `in_shape`: The shape of the input event.
+        - `out_shape`: The desired shape of the output event.
+        """
+        in_size = math.prod(in_shape)
+        out_size = math.prod(out_shape)
+
+        if in_size != out_size:
+            raise ValueError(
+                f"Shapes are incompatible: in_shape {in_shape} (size {in_size}) and "
+                f"out_shape {out_shape} (size {out_size}) must have the same total "
+                f"number of elements."
+            )
+
+        self.in_shape = in_shape
+        self.out_shape = out_shape
+
+    def forward_and_log_det(self, x: Array) -> tuple[Array, Array]:
+        """Computes y = reshape(x) and log|det J(f)(x)| = 0."""
+        y = jnp.reshape(x, self.out_shape)
+        return y, jnp.zeros((), dtype=x.dtype)
+
+    def inverse_and_log_det(self, y: Array) -> tuple[Array, Array]:
+        """Computes x = reshape(y) and log|det J(f^{-1})(y)| = 0."""
+        x = jnp.reshape(y, self.in_shape)
+        return x, jnp.zeros((), dtype=y.dtype)
+
+    def same_as(self, other: AbstractBijector) -> bool:
+        """Returns True if this bijector is guaranteed to be the same as `other`."""
+        return (
+            type(other) is Reshape
+            and self.in_shape == other.in_shape
+            and self.out_shape == other.out_shape
+        )

docs/api/bijectors/indexed.md

@@ -0,0 +1,7 @@
+# Indexed Bijector
+
+::: distreqx.bijectors.Indexed
+    options:
+        members:
+            - __init__
+---

docs/api/bijectors/permute.md

@@ -0,0 +1,7 @@
+# Permute Bijector
+
+::: distreqx.bijectors.Permute
+    options:
+        members:
+            - __init__
+---

docs/api/bijectors/reshape.md

@@ -0,0 +1,7 @@
+# Reshape Bijector
+
+::: distreqx.bijectors.Reshape
+    options:
+        members:
+            - __init__
+---

mkdocs.yml

@@ -124,6 +124,9 @@ nav:
             - 'api/bijectors/shift.md'
             - 'api/bijectors/sigmoid.md'
             - 'api/bijectors/tanh.md'
+            - 'api/bijectors/indexed.md'
+            - 'api/bijectors/permute.md'
+            - 'api/bijectors/reshape.md'
             - 'api/bijectors/triangular_linear.md'
             - 'api/bijectors/_bijector.md'
         - Utilities:

tests/conftest.py

@@ -0,0 +1,5 @@
+import jax
+
+# Must be set before any JAX arrays are initialized. Living here (rather than in
+# individual test files) guarantees it runs before any test module import.
+jax.config.update("jax_enable_x64", True)

tests/indexed_test.py

@@ -0,0 +1,59 @@
+from unittest import TestCase
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import numpy as np
+from parameterized import parameterized  # type: ignore
+
+from distreqx.bijectors import Indexed, ScalarAffine
+
+
+class IndexedTest(TestCase):
+    def setUp(self):
+        # We apply a ScalarAffine (y = 2x) only to indices 1 and 3
+        self.inner_bij = ScalarAffine(shift=jnp.array(0.0), scale=jnp.array(2.0))
+        self.bij = Indexed(bijector=self.inner_bij, indices=[1, 3])
+
+    def assertion_fn(self, rtol=1e-5):
+        return lambda x, y: np.testing.assert_allclose(x, y, rtol=rtol)
+
+    @parameterized.expand([("float32", jnp.float32), ("float64", jnp.float64)])
+    def test_forward_and_inverse(self, name, dtype):
+        x = jnp.array([10.0, 0.0, 20.0, 1.0], dtype=dtype)
+
+        y, log_det_fwd = self.bij.forward_and_log_det(x)
+
+        # Indices 0 and 2 unchanged. Indices 1 and 3 doubled.
+        expected_y = jnp.array([10.0, 0.0, 20.0, 2.0], dtype=dtype)
+        self.assertion_fn()(y, expected_y)
+
+        # The total log_det is the sum of the log_dets of the
+        # transformed indices (log(2) for each of index 1 and 3)
+        expected_logdet = jnp.array([jnp.log(2.0), jnp.log(2.0)], dtype=dtype)
+        self.assertion_fn()(log_det_fwd, expected_logdet)
+
+        x_rec, log_det_inv = self.bij.inverse_and_log_det(y)
+        self.assertion_fn()(x_rec, x)
+        self.assertion_fn()(log_det_inv, -expected_logdet)
+
+    def test_boolean_mask_indexing(self):
+        # Indexed can also accept boolean masks instead of integer indices
+        affine = ScalarAffine(shift=jnp.array(0.0), scale=jnp.array(2.0))
+        bool_bij = Indexed(bijector=affine, indices=[False, True, False, True])
+
+        x = jnp.array([10.0, 0.0, 20.0, 1.0])
+        y, log_det = bool_bij.forward_and_log_det(x)
+
+        expected_y = jnp.array([10.0, 0.0, 20.0, 2.0])
+        self.assertion_fn()(y, expected_y)
+
+    def test_jittable(self):
+        @eqx.filter_jit
+        def f(bij, x):
+            return bij.forward_and_log_det(x)
+
+        x = jnp.array([1.0, 2.0, 3.0, 4.0])
+        y, logdet = f(self.bij, x)
+        self.assertIsInstance(y, jax.Array)
+        self.assertIsInstance(logdet, jax.Array)

tests/permute_test.py

@@ -0,0 +1,50 @@
+from unittest import TestCase
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import numpy as np
+from parameterized import parameterized  # type: ignore
+
+from distreqx.bijectors import Permute
+
+
+class PermuteTest(TestCase):
+    def setUp(self):
+        self.bij = Permute(permutation=[2, 0, 1])
+
+    def assertion_fn(self, rtol=1e-5):
+        return lambda x, y: np.testing.assert_allclose(x, y, rtol=rtol)
+
+    def test_invalid_permutation(self):
+        with self.assertRaisesRegex(ValueError, "1D array"):
+            Permute([[0, 1], [2, 3]])
+
+        with self.assertRaisesRegex(ValueError, "exactly once"):
+            Permute([0, 1, 1])  # Duplicate 1, missing 2
+
+        with self.assertRaisesRegex(ValueError, "exactly once"):
+            Permute([0, 2, 3])  # Missing 1
+
+    @parameterized.expand([("float32", jnp.float32), ("float64", jnp.float64)])
+    def test_forward_and_inverse(self, name, dtype):
+        # Original: [A, B, C] -> Permuted: [C, A, B]
+        x = jnp.array([10.0, 20.0, 30.0], dtype=dtype)
+
+        y, log_det_fwd = self.bij.forward_and_log_det(x)
+        self.assertion_fn()(y, jnp.array([30.0, 10.0, 20.0], dtype=dtype))
+        self.assertEqual(log_det_fwd, 0.0)
+
+        x_rec, log_det_inv = self.bij.inverse_and_log_det(y)
+        self.assertion_fn()(x_rec, x)
+        self.assertEqual(log_det_inv, 0.0)
+
+    def test_jittable(self):
+        @eqx.filter_jit
+        def f(bij, x):
+            return bij.forward_and_log_det(x)
+
+        x = jnp.array([1.0, 2.0, 3.0])
+        y, logdet = f(self.bij, x)
+        self.assertIsInstance(y, jax.Array)
+        self.assertIsInstance(logdet, jax.Array)