Auto-differentiable numerical inverse & efficient materialization of BNAF masks

flowjax/bijections/utils.py

@@ -13,6 +13,7 @@
 from flowjax.bijections.chain import Chain
 from flowjax.utils import arraylike_to_array, check_shapes_match, merge_cond_shapes
 
+import lineax as lx
 
 class Invert(AbstractBijection):
     """Invert a bijection.
@@ -286,6 +287,9 @@ class NumericalInverse(AbstractBijection):
         bijection: The bijection to add an inverse to.
         inverter: Callable implementing the numerical inversion method. Should accept
             the bijection, y and condition as arguments, and return the inverse.
+        use_implicit_differentation: If ``True`` (default), gradients through the numerical
+            inverse are supplied with a custom JVP using implicit differentiation. If
+            ``False``, the inverter is assumed to be safely differentiable as provided.
     """
 
     bijection: AbstractBijection
@@ -297,13 +301,25 @@ def __init__(
         self,
         bijection: AbstractBijection,
         inverter: Callable[[AbstractBijection, Array, Array | None], Array],
+        use_implicit_differentation: bool = True,
     ):
+        self.bijection = bijection
+        self.shape = self.bijection.shape
+        self.cond_shape = self.bijection.cond_shape
+
+        if use_implicit_differentation:
+            self.inverter = self._wrap_inverter_with_implicit_jvp(inverter)
+        else:
+            self.inverter = inverter
+
+    @staticmethod
+    def _wrap_inverter_with_error_on_grad(inverter):
         @eqx.filter_custom_jvp
-        def nondiff_inverter(bijection, y, condition):
+        def wrapped_inverter(bijection, y, condition=None):
             return inverter(bijection, y, condition)
 
-        @nondiff_inverter.def_jvp
-        def nondiff_inverter_jvp(*args, **kwargs):
+        @wrapped_inverter.def_jvp
+        def wrapped_inverter_jvp(*_args, **_kwargs):
             raise RuntimeError(
                 "Computing gradients through the numerical inverse would lead to "
                 "misleading results. If you are using a flow with the analytical "
@@ -313,16 +329,46 @@ def nondiff_inverter_jvp(*args, **kwargs):
                 "supported)."
             )
 
-        self.bijection = bijection
-        self.inverter = nondiff_inverter
-        self.shape = self.bijection.shape
-        self.cond_shape = self.bijection.cond_shape
+        return wrapped_inverter
+
+    @staticmethod
+    def _wrap_inverter_with_implicit_jvp(inverter):
+        @eqx.filter_custom_jvp
+        def wrapped_inverter(bijection, y, condition=None):
+            return inverter(bijection, y, condition)
+
+        @wrapped_inverter.def_jvp
+        def wrapped_inverter_jvp(primals, tangents, condition=None):
+            (bijection, y), (bijection_tangent, y_tangent) = primals, tangents
+            x_star = wrapped_inverter(bijection, y, condition)
+            jacobian_operator = lx.JacobianLinearOperator(
+                lambda x, _: bijection.transform(x, condition),
+                x_star,
+            )
+
+            def residual(bijection, y):
+                return bijection.transform(x_star, condition) - y
+
+            _, residual_tangent = eqx.filter_jvp(
+                residual,
+                (bijection, y),
+                (bijection_tangent, y_tangent),
+            )
+
+            x_star_tangent = lx.linear_solve(
+                jacobian_operator,
+                -residual_tangent,
+            ).value
+
+            return x_star, x_star_tangent
+
+        return wrapped_inverter
 
     def transform_and_log_det(self, x, condition=None):
         return self.bijection.transform_and_log_det(x, condition)
 
     def inverse_and_log_det(self, y, condition=None):
-        x = self.inverter(self.bijection, y, condition)
+        x = self.inverter(self.bijection, y, condition=condition)
         _, log_det = self.bijection.transform_and_log_det(x, condition)
         return x, -log_det
 

flowjax/flows.py

@@ -168,6 +168,7 @@ def block_neural_autoregressive_flow(
     invert: bool = True,
     activation: AbstractBijection | Callable | None = None,
     inverter: Callable[[AbstractBijection, Array, Array | None], Array] | None = None,
+    use_implicit_differentation: bool = True
 ) -> Transformed:
     """Block neural autoregressive flow (BNAF) (https://arxiv.org/abs/1904.04676).
 
@@ -198,6 +199,9 @@ def block_neural_autoregressive_flow(
             invert the ``BlockAutoregressiveNetwork`` bijection. Passed to
             :py:class:`~flowjax.bijections.NumericalInverse`. Defaults to
             using ``elementwise_autoregressive_bisection``.
+        use_implicit_differentation:  Passed to :py:class:`~flowjax.bijections.NumericalInverse`.
+            If ``True`` (default), gradients through the through the inverse use a custom JVP. 
+            If ``False``, the inverter is treated as differentiable.
     """
     dim = base_dist.shape[-1]
 
@@ -222,6 +226,7 @@ def make_layer(key):  # bnaf layer + permutation
                 activation=activation,
             ),
             inverter=inverter,
+            use_implicit_differentation=use_implicit_differentation
         )
         return _add_default_permute(bijection, base_dist.shape[-1], perm_key)
 

flowjax/masks.py

@@ -8,7 +8,6 @@
 import operator
 
 import jax.numpy as jnp
-from jax.scipy.linalg import block_diag
 from jaxtyping import Array, Bool, Int
 
 
@@ -37,16 +36,15 @@ def rank_based_mask(
 
 def block_diag_mask(block_shape: tuple, n_blocks: int) -> Bool[Array, "dim1 dim2"]:
     """Block diagonal mask."""
-    return block_diag(*jnp.ones((n_blocks, *block_shape), bool))
+    mask = jnp.diag(jnp.ones(n_blocks, dtype=int))
+    block = jnp.ones(block_shape, dtype=int)
+    return jnp.kron(mask, block).astype(bool)
 
 
 def block_tril_mask(
     block_shape: tuple, n_blocks: int, k: int = 0
 ) -> Bool[Array, "dim1 dim2"]:
     """Lower triangular block mask, with offset k."""
-    mask = jnp.zeros((block_shape[0] * n_blocks, block_shape[1] * n_blocks), bool)
-    for i in range(n_blocks):
-        row_i = max(0, (i - k)) * block_shape[0]
-        col_i = i * block_shape[1]
-        mask = mask.at[row_i:, col_i : col_i + block_shape[1]].set(True)  # noqa
-    return mask
+    mask = jnp.tri(n_blocks, n_blocks, k=k, dtype=int)
+    block = jnp.ones(block_shape, dtype=int)
+    return jnp.kron(mask, block).astype(bool)

pyproject.toml

@@ -16,6 +16,7 @@ dependencies = [
     "tqdm",
     "optax",
     "paramax>=0.0.5",
+    "lineax"
 ]
 description = "Easy to use distributions, bijections and normalizing flows in JAX."
 keywords = ["jax", "neural-networks", "equinox"]

tests/test_bijections/test_bijection_utils.py

@@ -42,14 +42,12 @@ def test_Permute_argcheck():
 
 
 @pytest.mark.parametrize("diff_fn", test_cases)
-def test_NumericalInverse_not_differentiable(diff_fn):
+def test_NumericalInverse_implicit_jvp(diff_fn):
     bijection = NumericalInverse(
         Affine(5, 2),
         root_finder_to_inverter(
             partial(bisection_search, lower=-1, upper=1, atol=1e-7),
         ),
     )
-    with pytest.raises(
-        RuntimeError, match="Computing gradients through the numerical inverse"
-    ):
-        diff_fn(bijection.inverse)(jnp.ones(()))
+    derivative = diff_fn(bijection.inverse)(jnp.array(1.0))
+    assert jnp.allclose(derivative, 0.5, atol=1e-5)

tests/test_flows.py

@@ -22,6 +22,7 @@
 
 testcases = {
     "BNAF": lambda c: block_neural_autoregressive_flow(**KWARGS, cond_dim=c),
+    "BNAF_diffable_inverse": lambda c: block_neural_autoregressive_flow(**KWARGS, invert=False, cond_dim=c),
     "triangular_spline_flow": lambda c: triangular_spline_flow(**KWARGS, cond_dim=c),
     "affine_coupling": lambda c: coupling_flow(
         transformer=Affine(), **KWARGS, cond_dim=c