diff --git a/distreqx/bijectors/__init__.py b/distreqx/bijectors/__init__.py index 7805142..a09dee6 100644 --- a/distreqx/bijectors/__init__.py +++ b/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 diff --git a/distreqx/bijectors/_restructure.py b/distreqx/bijectors/_restructure.py new file mode 100644 index 0000000..8c2b3e3 --- /dev/null +++ b/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 + ) diff --git a/docs/api/bijectors/restructure.md b/docs/api/bijectors/restructure.md new file mode 100644 index 0000000..36b8202 --- /dev/null +++ b/docs/api/bijectors/restructure.md @@ -0,0 +1,7 @@ +# Restructure Bijector + +::: distreqx.bijectors.Restructure + options: + members: + - __init__ +--- \ No newline at end of file diff --git a/mkdocs.yml b/mkdocs.yml index aa5cf4e..a6778ed 100644 --- a/mkdocs.yml +++ b/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: diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..8f407ab --- /dev/null +++ b/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) diff --git a/tests/restructure_test.py b/tests/restructure_test.py new file mode 100644 index 0000000..d4bcd01 --- /dev/null +++ b/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))