Add Restructure

distreqx/bijectors/__init__.py

@@ -8,6 +8,7 @@
 from ._chain import Chain as Chain
 from ._diag_linear import DiagLinear as DiagLinear
 from ._linear import AbstractLinearBijector as AbstractLinearBijector
+from ._restructure import Restructure as Restructure
 from ._scalar_affine import ScalarAffine as ScalarAffine
 from ._shift import Shift as Shift
 from ._sigmoid import Sigmoid as Sigmoid

distreqx/bijectors/_restructure.py

@@ -0,0 +1,102 @@
+from typing import Any
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+from jaxtyping import Array, PyTreeDef
+
+from ._bijector import (
+    AbstractBijector,
+    AbstractForwardInverseBijector,
+    AbstractFwdLogDetJacBijector,
+    AbstractInvLogDetJacBijector,
+)
+
+
+class Restructure(
+    AbstractForwardInverseBijector,
+    AbstractInvLogDetJacBijector,
+    AbstractFwdLogDetJacBijector,
+):
+    """A bijector that restructures a PyTree of arrays.
+
+    This is equivalent to `tfp.bijectors.Restructure`. It maps values between
+    different nested structures (e.g., lists to dicts) without modifying the
+    underlying arrays themselves.
+    """
+
+    in_treedef: PyTreeDef = eqx.field(static=True)  # type: ignore
+    out_treedef: PyTreeDef = eqx.field(static=True)  # type: ignore
+    forward_permutation: tuple[int, ...] = eqx.field(static=True)
+    inverse_permutation: tuple[int, ...] = eqx.field(static=True)
+
+    _is_constant_jacobian: bool = True
+    _is_constant_log_det: bool = True
+
+    def __init__(self, in_structure: Any, out_structure: Any):
+        """Initializes a Restructure bijector.
+
+        **Arguments:**
+
+        - `in_structure`: A PyTree defining the input structure. Its leaves must
+          be unique identifier tokens (e.g., integers or strings).
+        - `out_structure`: A PyTree defining the desired output structure. It must
+          contain the exact same set of tokens as `in_structure`.
+        """
+        self.in_treedef = jax.tree_util.tree_structure(in_structure)
+        self.out_treedef = jax.tree_util.tree_structure(out_structure)
+
+        flat_in = jax.tree_util.tree_leaves(in_structure)
+        flat_out = jax.tree_util.tree_leaves(out_structure)
+
+        if len(flat_in) != len(set(flat_in)):
+            raise ValueError(
+                f"in_structure cannot have duplicate tokens. Got: {flat_in}"
+            )
+        if len(flat_out) != len(set(flat_out)):
+            raise ValueError(
+                f"out_structure cannot have duplicate tokens. Got: {flat_out}"
+            )
+        if set(flat_in) != set(flat_out):
+            raise ValueError(
+                f"Structures are incompatible: in_structure tokens {set(flat_in)} "
+                f"do not match out_structure tokens {set(flat_out)}."
+            )
+
+        # Pre-compute the routing permutations for fast forward/inverse passes
+        self.forward_permutation = tuple(flat_in.index(token) for token in flat_out)
+        self.inverse_permutation = tuple(flat_out.index(token) for token in flat_in)
+
+    def forward_and_log_det(self, x: Any) -> tuple[Any, Array]:
+        """Computes y = restructure(x) and log|det J(f)(x)| = 0."""
+        if jax.tree_util.tree_structure(x) != self.in_treedef:
+            raise ValueError("Input `x` does not match the expected `in_structure`.")
+
+        flat_x = jax.tree_util.tree_leaves(x)
+        flat_y = [flat_x[i] for i in self.forward_permutation]
+        y = jax.tree_util.tree_unflatten(self.out_treedef, flat_y)
+
+        # Pull dtype from the first leaf to match JAX array constraints
+        dtype = flat_x[0].dtype if flat_x else jnp.float32
+        return y, jnp.zeros((), dtype=dtype)
+
+    def inverse_and_log_det(self, y: Any) -> tuple[Any, Array]:
+        """Computes x = restructure^{-1}(y) and log|det J(f^{-1})(y)| = 0."""
+        if jax.tree_util.tree_structure(y) != self.out_treedef:
+            raise ValueError("Input `y` does not match the expected `out_structure`.")
+
+        flat_y = jax.tree_util.tree_leaves(y)
+        flat_x = [flat_y[i] for i in self.inverse_permutation]
+        x = jax.tree_util.tree_unflatten(self.in_treedef, flat_x)
+
+        dtype = flat_y[0].dtype if flat_y else jnp.float32
+        return x, jnp.zeros((), dtype=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 Restructure
+            and self.in_treedef == other.in_treedef
+            and self.out_treedef == other.out_treedef
+            and self.forward_permutation == other.forward_permutation
+        )

docs/api/bijectors/restructure.md

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

mkdocs.yml

@@ -124,6 +124,7 @@ nav:
             - 'api/bijectors/shift.md'
             - 'api/bijectors/sigmoid.md'
             - 'api/bijectors/tanh.md'
+            - 'api/bijectors/restructure.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/restructure_test.py

@@ -0,0 +1,122 @@
+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 Restructure
+
+
+class RestructureTest(TestCase):
+    def setUp(self):
+        # We will test mapping a 3-element list to a dictionary
+        self.in_structure = [0, 1, 2]
+        self.out_structure = {"a": 0, "b": 2, "c": 1}
+        self.bij = Restructure(
+            in_structure=self.in_structure, out_structure=self.out_structure
+        )
+
+    def assertion_fn(self, rtol=1e-5):
+        return lambda x, y: np.testing.assert_allclose(x, y, rtol=rtol)
+
+    def test_invalid_structures(self):
+        # Duplicate tokens in the input structure
+        with self.assertRaisesRegex(ValueError, "duplicate tokens"):
+            Restructure([0, 0, 1], {"a": 0, "b": 1, "c": 2})
+
+        # Duplicate tokens in the output structure
+        with self.assertRaisesRegex(ValueError, "duplicate tokens"):
+            Restructure([0, 1, 2], {"a": 0, "b": 0, "c": 1})
+
+        # Mismatched token sets
+        with self.assertRaisesRegex(ValueError, "incompatible"):
+            Restructure([0, 1, 2], {"a": 0, "b": 1, "c": 3})
+
+    def test_invalid_input_trees(self):
+        # Passing the wrong structure to forward
+        # (e.g., passing the dict instead of the list)
+        bad_x = {"a": jnp.ones(2), "b": jnp.ones(2), "c": jnp.ones(2)}
+        with self.assertRaisesRegex(ValueError, "does not match the expected"):
+            self.bij.forward_and_log_det(bad_x)
+
+        # Passing the wrong structure to inverse
+        bad_y = [jnp.ones(2), jnp.ones(2), jnp.ones(2)]
+        with self.assertRaisesRegex(ValueError, "does not match the expected"):
+            self.bij.inverse_and_log_det(bad_y)
+
+    @parameterized.expand([("float32", jnp.float32), ("float64", jnp.float64)])
+    def test_forward_and_log_det(self, name, dtype):
+        x = [
+            jnp.array([1.0, 2.0], dtype=dtype),
+            jnp.array([[3.0]], dtype=dtype),
+            jnp.array(4.0, dtype=dtype),
+        ]
+        y, log_det = self.bij.forward_and_log_det(x)
+
+        # Verify the structure and values map correctly according to the tokens
+        self.assertIsInstance(y, dict)
+        self.assertEqual(set(y.keys()), {"a", "b", "c"})
+        self.assertion_fn()(y["a"], x[0])  # Token 0
+        self.assertion_fn()(y["c"], x[1])  # Token 1
+        self.assertion_fn()(y["b"], x[2])  # Token 2
+
+        # log_det must be an unbatched scalar 0.0 of the matching dtype
+        self.assertEqual(log_det.shape, ())
+        self.assertEqual(log_det, 0.0)
+
+        # Check dtype preservation
+        self.assertEqual(y["a"].dtype, dtype)
+        self.assertEqual(log_det.dtype, dtype)
+
+    @parameterized.expand([("float32", jnp.float32), ("float64", jnp.float64)])
+    def test_inverse_and_log_det(self, name, dtype):
+        y = {
+            "a": jnp.array([1.0, 2.0], dtype=dtype),
+            "b": jnp.array(4.0, dtype=dtype),
+            "c": jnp.array([[3.0]], dtype=dtype),
+        }
+        x, log_det = self.bij.inverse_and_log_det(y)
+
+        # Verify the structure and values map correctly back to the list
+        self.assertIsInstance(x, list)
+        self.assertEqual(len(x), 3)
+        self.assertion_fn()(x[0], y["a"])  # Token 0
+        self.assertion_fn()(x[1], y["c"])  # Token 1
+        self.assertion_fn()(x[2], y["b"])  # Token 2
+
+        self.assertEqual(log_det.shape, ())
+        self.assertEqual(log_det, 0.0)
+        self.assertEqual(x[0].dtype, dtype)
+        self.assertEqual(log_det.dtype, dtype)
+
+    def test_jittable(self):
+        @eqx.filter_jit
+        def f_forward(bij, x):
+            return bij.forward_and_log_det(x)
+
+        @eqx.filter_jit
+        def f_inverse(bij, y):
+            return bij.inverse_and_log_det(y)
+
+        x = [jnp.ones(1), jnp.ones(2), jnp.ones(3)]
+        y, log_det_fwd = f_forward(self.bij, x)
+
+        self.assertIsInstance(y, dict)
+        self.assertIsInstance(log_det_fwd, jax.Array)
+
+        x_reconstructed, log_det_inv = f_inverse(self.bij, y)
+        self.assertIsInstance(x_reconstructed, list)
+        self.assertIsInstance(log_det_inv, jax.Array)
+
+    def test_same_as(self):
+        same_bij = Restructure(
+            in_structure=[0, 1, 2], out_structure={"a": 0, "b": 2, "c": 1}
+        )
+        diff_bij = Restructure(
+            in_structure=[0, 1, 2], out_structure={"a": 1, "b": 2, "c": 0}
+        )
+
+        self.assertTrue(self.bij.same_as(same_bij))
+        self.assertFalse(self.bij.same_as(diff_bij))