From fbc9f3bd1b3d275f6f5209e4e7ff171c0415a9c0 Mon Sep 17 00:00:00 2001
From: Noah Wolfe <newolfe@mit.edu>
Date: Mon, 18 Aug 2025 10:23:15 -0400
Subject: [PATCH 01/10] Modify pass through arguments for use with optimistix;
 allow switch to turn off custom jvp warning preventing autodiff of inverter

---
 flowjax/bijections/utils.py | 35 ++++++++++++++++++++---------------
 flowjax/flows.py            |  2 ++
 flowjax/root_finding.py     |  2 +-
 3 files changed, 23 insertions(+), 16 deletions(-)

diff --git a/flowjax/bijections/utils.py b/flowjax/bijections/utils.py
index a75333b5..34995625 100644
--- a/flowjax/bijections/utils.py
+++ b/flowjax/bijections/utils.py
@@ -292,24 +292,29 @@ def __init__(
         self,
         bijection: AbstractBijection,
         inverter: Callable[[AbstractBijection, Array, Array | None], Array],
+        diffable_inverter: bool = False
     ):
-        @eqx.filter_custom_jvp
-        def nondiff_inverter(bijection, y, condition):
-            return inverter(bijection, y, condition)
-
-        @nondiff_inverter.def_jvp
-        def nondiff_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 "
-                "transform only defined in one direction, consider inverting the "
-                "bijection by flipping the ``invert`` argument to the flow. If this is "
-                "not possible, consider using implicit differentation (not yet "
-                "supported)."
-            )
+        if not diffable_inverter:
+            @eqx.filter_custom_jvp
+            def nondiff_inverter(bijection, y, condition):
+                return inverter(bijection, y, condition)
+
+            @nondiff_inverter.def_jvp
+            def nondiff_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 "
+                    "transform only defined in one direction, consider inverting the "
+                    "bijection by flipping the ``invert`` argument to the flow. If this is "
+                    "not possible, consider using implicit differentation (not yet "
+                    "supported)."
+                )
+
+            self.inverter = nondiff_inverter
+        else:
+            self.inverter = inverter
 
         self.bijection = bijection
-        self.inverter = nondiff_inverter
         self.shape = self.bijection.shape
         self.cond_shape = self.bijection.cond_shape
 
diff --git a/flowjax/flows.py b/flowjax/flows.py
index 231b310e..4400dc2a 100644
--- a/flowjax/flows.py
+++ b/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,
+    diffable_inverter: bool = False
 ) -> Transformed:
     """Block neural autoregressive flow (BNAF) (https://arxiv.org/abs/1904.04676).
 
@@ -222,6 +223,7 @@ def make_layer(key):  # bnaf layer + permutation
                 activation=activation,
             ),
             inverter=inverter,
+            diffable_inverter=diffable_inverter
         )
         return _add_default_permute(bijection, base_dist.shape[-1], perm_key)
 
diff --git a/flowjax/root_finding.py b/flowjax/root_finding.py
index 468969b7..52a141d9 100644
--- a/flowjax/root_finding.py
+++ b/flowjax/root_finding.py
@@ -85,7 +85,7 @@ def inverter(bijection, y, condition=None):
         def fn(x):
             return bijection.transform(x, condition) - y
 
-        return root_finder(fn)[0]
+        return root_finder(fn, y)[0]
 
     return inverter
 

From ce942933c18be2eba9951ddc4b948e32f6373239 Mon Sep 17 00:00:00 2001
From: Noah Wolfe <newolfe@mit.edu>
Date: Tue, 26 Aug 2025 16:09:27 -0400
Subject: [PATCH 02/10] custom jvp using implicit differentation to compute
 derivaties of the inverse (numerical) transformation

---
 flowjax/bijections/utils.py | 35 +++++++++++++++++++++++------------
 flowjax/root_finding.py     |  2 +-
 2 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/flowjax/bijections/utils.py b/flowjax/bijections/utils.py
index 34995625..d7e85241 100644
--- a/flowjax/bijections/utils.py
+++ b/flowjax/bijections/utils.py
@@ -9,6 +9,8 @@
 import numpy as np
 from jaxtyping import Array, Int
 
+from jax import jacfwd
+
 from flowjax.bijections.bijection import AbstractBijection
 from flowjax.bijections.chain import Chain
 from flowjax.utils import arraylike_to_array, check_shapes_match, merge_cond_shapes
@@ -296,21 +298,30 @@ def __init__(
     ):
         if not diffable_inverter:
             @eqx.filter_custom_jvp
-            def nondiff_inverter(bijection, y, condition):
+            def inverter_wrapper(bijection, y, condition=None):
                 return inverter(bijection, y, condition)
 
-            @nondiff_inverter.def_jvp
-            def nondiff_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 "
-                    "transform only defined in one direction, consider inverting the "
-                    "bijection by flipping the ``invert`` argument to the flow. If this is "
-                    "not possible, consider using implicit differentation (not yet "
-                    "supported)."
+            @inverter_wrapper.def_jvp
+            def inverter_jvp(primals, tangents, condition=None):
+                (bijection, y), (bijection_dot, y_dot) = primals, tangents
+
+                x_star = inverter_wrapper(bijection, y, condition)
+
+                # TODO: implement with lineax JacobianLinearOperator
+                # which determines whether jacfwd or jacrev is more efficient
+                A = jacfwd(lambda x: bijection.transform(x, condition))(x_star)
+
+                def F(bijection, y):
+                    return bijection.transform(x_star, condition) - y
+
+                _, b = eqx.filter_jvp(
+                    F, (bijection, y), (bijection_dot, y_dot)
                 )
 
-            self.inverter = nondiff_inverter
+                # TODO: use more robust solvers in lineax
+                return x_star, jnp.linalg.solve(A, b)
+
+            self.inverter = inverter_wrapper
         else:
             self.inverter = inverter
 
@@ -322,7 +333,7 @@ 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
 
diff --git a/flowjax/root_finding.py b/flowjax/root_finding.py
index 52a141d9..468969b7 100644
--- a/flowjax/root_finding.py
+++ b/flowjax/root_finding.py
@@ -85,7 +85,7 @@ def inverter(bijection, y, condition=None):
         def fn(x):
             return bijection.transform(x, condition) - y
 
-        return root_finder(fn, y)[0]
+        return root_finder(fn)[0]
 
     return inverter
 

From f1cd4156628b0a32353de217a283c33fdbe300ef Mon Sep 17 00:00:00 2001
From: Noah Wolfe <newolfe@mit.edu>
Date: Tue, 26 Aug 2025 16:58:00 -0400
Subject: [PATCH 03/10] fix sign error in implicit diff and implement with
 lineax

---
 flowjax/bijections/utils.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/flowjax/bijections/utils.py b/flowjax/bijections/utils.py
index d7e85241..f12702d4 100644
--- a/flowjax/bijections/utils.py
+++ b/flowjax/bijections/utils.py
@@ -15,6 +15,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.
@@ -309,7 +310,11 @@ def inverter_jvp(primals, tangents, condition=None):
 
                 # TODO: implement with lineax JacobianLinearOperator
                 # which determines whether jacfwd or jacrev is more efficient
-                A = jacfwd(lambda x: bijection.transform(x, condition))(x_star)
+                #A = jacfwd(lambda x: bijection.transform(x, condition))(x_star)
+                A = lx.JacobianLinearOperator(
+                    lambda x, _: bijection.transform(x, condition),
+                    x_star
+                )
 
                 def F(bijection, y):
                     return bijection.transform(x_star, condition) - y
@@ -319,7 +324,7 @@ def F(bijection, y):
                 )
 
                 # TODO: use more robust solvers in lineax
-                return x_star, jnp.linalg.solve(A, b)
+                return x_star, lx.linear_solve(A, -b).value #jnp.linalg.solve(A, b)
 
             self.inverter = inverter_wrapper
         else:

From 0d55c67ada8fd095b47b7dffda33affc82006d8a Mon Sep 17 00:00:00 2001
From: Noah Wolfe <newolfe@mit.edu>
Date: Fri, 19 Sep 2025 11:04:29 -0400
Subject: [PATCH 04/10] push latest local copy of working invertible bnaf

---
 .../block_autoregressive_network.py           |  4 +
 flowjax/bijections/utils.py                   | 73 ++++++++++++-------
 flowjax/flows.py                              |  6 +-
 3 files changed, 54 insertions(+), 29 deletions(-)

diff --git a/flowjax/bijections/block_autoregressive_network.py b/flowjax/bijections/block_autoregressive_network.py
index 8a15c8e5..e1484de5 100644
--- a/flowjax/bijections/block_autoregressive_network.py
+++ b/flowjax/bijections/block_autoregressive_network.py
@@ -122,6 +122,10 @@ def __init__(
                 (1, block_dim),
             ]
 
+            def make_layer(inp):
+                layer_key, block_shape = inp
+                return block_autoregressive_linear(layer_key, n_blocks=dim, block_shape=block_shape)
+
             for layer_key, block_shape in zip(keys, block_shapes, strict=True):
                 layers_and_log_jac_fns.append(
                     block_autoregressive_linear(
diff --git a/flowjax/bijections/utils.py b/flowjax/bijections/utils.py
index f12702d4..e9edae11 100644
--- a/flowjax/bijections/utils.py
+++ b/flowjax/bijections/utils.py
@@ -295,40 +295,59 @@ def __init__(
         self,
         bijection: AbstractBijection,
         inverter: Callable[[AbstractBijection, Array, Array | None], Array],
-        diffable_inverter: bool = False
+        diffable_inverter: bool = False,
+        raise_old_error: bool = False
     ):
-        if not diffable_inverter:
+        if raise_old_error:
             @eqx.filter_custom_jvp
-            def inverter_wrapper(bijection, y, condition=None):
+            def nondiff_inverter(bijection, y, condition):
                 return inverter(bijection, y, condition)
 
-            @inverter_wrapper.def_jvp
-            def inverter_jvp(primals, tangents, condition=None):
-                (bijection, y), (bijection_dot, y_dot) = primals, tangents
-
-                x_star = inverter_wrapper(bijection, y, condition)
-
-                # TODO: implement with lineax JacobianLinearOperator
-                # which determines whether jacfwd or jacrev is more efficient
-                #A = jacfwd(lambda x: bijection.transform(x, condition))(x_star)
-                A = lx.JacobianLinearOperator(
-                    lambda x, _: bijection.transform(x, condition),
-                    x_star
-                )
-
-                def F(bijection, y):
-                    return bijection.transform(x_star, condition) - y
-
-                _, b = eqx.filter_jvp(
-                    F, (bijection, y), (bijection_dot, y_dot)
+            @nondiff_inverter.def_jvp
+            def nondiff_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 "
+                    "transform only defined in one direction, consider inverting the "
+                    "bijection by flipping the ``invert`` argument to the flow. If this is "
+                    "not possible, consider using implicit differentation (not yet "
+                    "supported)."
                 )
 
-                # TODO: use more robust solvers in lineax
-                return x_star, lx.linear_solve(A, -b).value #jnp.linalg.solve(A, b)
-
-            self.inverter = inverter_wrapper
+            self.inverter = nondiff_inverter
         else:
-            self.inverter = inverter
+            if not diffable_inverter:
+                @eqx.filter_custom_jvp
+                def inverter_wrapper(bijection, y, condition=None):
+                    return inverter(bijection, y, condition)
+
+                @inverter_wrapper.def_jvp
+                def inverter_jvp(primals, tangents, condition=None):
+                    (bijection, y), (bijection_dot, y_dot) = primals, tangents
+
+                    x_star = inverter_wrapper(bijection, y, condition)
+
+                    # TODO: implement with lineax JacobianLinearOperator
+                    # which determines whether jacfwd or jacrev is more efficient
+                    #A = jacfwd(lambda x: bijection.transform(x, condition))(x_star)
+                    A = lx.JacobianLinearOperator(
+                        lambda x, _: bijection.transform(x, condition),
+                        x_star
+                    )
+
+                    def F(bijection, y):
+                        return bijection.transform(x_star, condition) - y
+
+                    _, b = eqx.filter_jvp(
+                        F, (bijection, y), (bijection_dot, y_dot)
+                    )
+
+                    # TODO: use more robust solvers in lineax
+                    return x_star, lx.linear_solve(A, -b).value #jnp.linalg.solve(A, b)
+
+                self.inverter = inverter_wrapper
+            else:
+                self.inverter = inverter
 
         self.bijection = bijection
         self.shape = self.bijection.shape
diff --git a/flowjax/flows.py b/flowjax/flows.py
index 4400dc2a..a645a662 100644
--- a/flowjax/flows.py
+++ b/flowjax/flows.py
@@ -168,7 +168,8 @@ def block_neural_autoregressive_flow(
     invert: bool = True,
     activation: AbstractBijection | Callable | None = None,
     inverter: Callable[[AbstractBijection, Array, Array | None], Array] | None = None,
-    diffable_inverter: bool = False
+    diffable_inverter: bool = False,
+    raise_old_error: bool = False
 ) -> Transformed:
     """Block neural autoregressive flow (BNAF) (https://arxiv.org/abs/1904.04676).
 
@@ -223,7 +224,8 @@ def make_layer(key):  # bnaf layer + permutation
                 activation=activation,
             ),
             inverter=inverter,
-            diffable_inverter=diffable_inverter
+            diffable_inverter=diffable_inverter,
+            raise_old_error=raise_old_error
         )
         return _add_default_permute(bijection, base_dist.shape[-1], perm_key)
 

From 1dcebd7ddc78b6f2fb5299e4a6adb270aa67263a Mon Sep 17 00:00:00 2001
From: Noah Wolfe <newolfe@mit.edu>
Date: Wed, 24 Sep 2025 13:29:05 -0400
Subject: [PATCH 05/10] efficient block-lower triangular mask allocation

---
 flowjax/masks.py | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/flowjax/masks.py b/flowjax/masks.py
index 71a1ced1..1de2f162 100644
--- a/flowjax/masks.py
+++ b/flowjax/masks.py
@@ -44,9 +44,7 @@ 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
+    mask = jnp.tri(n_blocks, n_blocks, k=k, dtype=int)
+    block = jnp.ones(block_shape, dtype=int)
+    mask = jnp.kron(mask, block).astype(bool)
     return mask

From feb834b902a980ce75801253e3f2ab94f896aaf2 Mon Sep 17 00:00:00 2001
From: Noah Wolfe <newolfe@mit.edu>
Date: Wed, 24 Sep 2025 14:37:38 -0400
Subject: [PATCH 06/10] efficient block-diagonal mask allocation

---
 flowjax/masks.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/flowjax/masks.py b/flowjax/masks.py
index 1de2f162..833052b3 100644
--- a/flowjax/masks.py
+++ b/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,7 +36,9 @@ 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(
@@ -46,5 +47,4 @@ def block_tril_mask(
     """Lower triangular block mask, with offset k."""
     mask = jnp.tri(n_blocks, n_blocks, k=k, dtype=int)
     block = jnp.ones(block_shape, dtype=int)
-    mask = jnp.kron(mask, block).astype(bool)
-    return mask
+    return jnp.kron(mask, block).astype(bool)

From 8f8b7e36aa3d906653ce6d56ac6664ce7f94a130 Mon Sep 17 00:00:00 2001
From: Noah Wolfe <newolfe@mit.edu>
Date: Mon, 9 Feb 2026 16:21:18 -0500
Subject: [PATCH 07/10] cleanup to code, documentation, and add light unit
 tests

---
 flowjax/bijections/utils.py                   | 109 ++++++++++--------
 flowjax/flows.py                              |   6 +
 tests/test_bijections/test_bijection_utils.py |  16 ++-
 tests/test_flows.py                           |   1 +
 4 files changed, 82 insertions(+), 50 deletions(-)

diff --git a/flowjax/bijections/utils.py b/flowjax/bijections/utils.py
index e9edae11..32e5bcc1 100644
--- a/flowjax/bijections/utils.py
+++ b/flowjax/bijections/utils.py
@@ -9,8 +9,6 @@
 import numpy as np
 from jaxtyping import Array, Int
 
-from jax import jacfwd
-
 from flowjax.bijections.bijection import AbstractBijection
 from flowjax.bijections.chain import Chain
 from flowjax.utils import arraylike_to_array, check_shapes_match, merge_cond_shapes
@@ -284,6 +282,12 @@ 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.
+        diffable_inverter: If ``False`` (default), gradients through the numerical
+            inverse are supplied with a custom JVP using implicit differentiation. If
+            ``True``, the inverter is assumed to be safely differentiable as provided.
+        raise_old_error: If ``True``, preserve legacy behaviour and raise
+            ``RuntimeError`` when users attempt to differentiate through the numerical
+            inverse. Takes precedence over ``diffable_inverter``.
     """
 
     bijection: AbstractBijection
@@ -296,62 +300,69 @@ def __init__(
         bijection: AbstractBijection,
         inverter: Callable[[AbstractBijection, Array, Array | None], Array],
         diffable_inverter: bool = False,
-        raise_old_error: bool = False
+        raise_old_error: bool = False,
     ):
+        self.bijection = bijection
+        self.shape = self.bijection.shape
+        self.cond_shape = self.bijection.cond_shape
         if raise_old_error:
-            @eqx.filter_custom_jvp
-            def nondiff_inverter(bijection, y, condition):
-                return inverter(bijection, y, condition)
-
-            @nondiff_inverter.def_jvp
-            def nondiff_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 "
-                    "transform only defined in one direction, consider inverting the "
-                    "bijection by flipping the ``invert`` argument to the flow. If this is "
-                    "not possible, consider using implicit differentation (not yet "
-                    "supported)."
-                )
-
-            self.inverter = nondiff_inverter
+            self.inverter = self._wrap_inverter_with_error_on_grad(inverter)
+        elif diffable_inverter:
+            self.inverter = inverter
         else:
-            if not diffable_inverter:
-                @eqx.filter_custom_jvp
-                def inverter_wrapper(bijection, y, condition=None):
-                    return inverter(bijection, y, condition)
-
-                @inverter_wrapper.def_jvp
-                def inverter_jvp(primals, tangents, condition=None):
-                    (bijection, y), (bijection_dot, y_dot) = primals, tangents
-
-                    x_star = inverter_wrapper(bijection, y, condition)
+            self.inverter = self._wrap_inverter_with_implicit_jvp(inverter)
+
+    @staticmethod
+    def _wrap_inverter_with_error_on_grad(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(*_args, **_kwargs):
+            raise RuntimeError(
+                "Computing gradients through the numerical inverse would lead to "
+                "misleading results. If you are using a flow with the analytical "
+                "transform only defined in one direction, consider inverting the "
+                "bijection by flipping the ``invert`` argument to the flow. If this is "
+                "not possible, consider using implicit differentation (not yet "
+                "supported)."
+            )
 
-                    # TODO: implement with lineax JacobianLinearOperator
-                    # which determines whether jacfwd or jacrev is more efficient
-                    #A = jacfwd(lambda x: bijection.transform(x, condition))(x_star)
-                    A = lx.JacobianLinearOperator(
-                        lambda x, _: bijection.transform(x, condition),
-                        x_star
-                    )
+        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 F(bijection, y):
-                        return bijection.transform(x_star, condition) - y
+            def residual(bijection, y):
+                return bijection.transform(x_star, condition) - y
 
-                    _, b = eqx.filter_jvp(
-                        F, (bijection, y), (bijection_dot, y_dot)
-                    )
+            _, residual_tangent = eqx.filter_jvp(
+                residual,
+                (bijection, y),
+                (bijection_tangent, y_tangent),
+            )
 
-                    # TODO: use more robust solvers in lineax
-                    return x_star, lx.linear_solve(A, -b).value #jnp.linalg.solve(A, b)
+            x_star_tangent = lx.linear_solve(
+                jacobian_operator,
+                -residual_tangent,
+            ).value
 
-                self.inverter = inverter_wrapper
-            else:
-                self.inverter = inverter
+            return x_star, x_star_tangent
 
-        self.bijection = bijection
-        self.shape = self.bijection.shape
-        self.cond_shape = self.bijection.cond_shape
+        return wrapped_inverter
 
     def transform_and_log_det(self, x, condition=None):
         return self.bijection.transform_and_log_det(x, condition)
diff --git a/flowjax/flows.py b/flowjax/flows.py
index a645a662..cc01658b 100644
--- a/flowjax/flows.py
+++ b/flowjax/flows.py
@@ -200,6 +200,12 @@ def block_neural_autoregressive_flow(
             invert the ``BlockAutoregressiveNetwork`` bijection. Passed to
             :py:class:`~flowjax.bijections.NumericalInverse`. Defaults to
             using ``elementwise_autoregressive_bisection``.
+        diffable_inverter: Passed to :py:class:`~flowjax.bijections.NumericalInverse`.
+            If ``False`` (default), gradients through the inverse use a custom JVP.
+            If ``True``, the provided ``inverter`` is treated as differentiable.
+        raise_old_error: Passed to :py:class:`~flowjax.bijections.NumericalInverse`.
+            If ``True``, attempting to differentiate through the numerical inverse
+            raises a ``RuntimeError`` (legacy behavior).
     """
     dim = base_dist.shape[-1]
 
diff --git a/tests/test_bijections/test_bijection_utils.py b/tests/test_bijections/test_bijection_utils.py
index f03540c8..22316f80 100644
--- a/tests/test_bijections/test_bijection_utils.py
+++ b/tests/test_bijections/test_bijection_utils.py
@@ -42,13 +42,27 @@ 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),
         ),
     )
+    derivative = diff_fn(bijection.inverse)(jnp.array(1.0))
+    assert jnp.allclose(derivative, 0.5, atol=1e-5)
+
+
+@pytest.mark.parametrize("diff_fn", test_cases)
+def test_NumericalInverse_raise_old_error_takes_precedence(diff_fn):
+    bijection = NumericalInverse(
+        Affine(5, 2),
+        root_finder_to_inverter(
+            partial(bisection_search, lower=-1, upper=1, atol=1e-7),
+        ),
+        diffable_inverter=True,
+        raise_old_error=True,
+    )
     with pytest.raises(
         RuntimeError, match="Computing gradients through the numerical inverse"
     ):
diff --git a/tests/test_flows.py b/tests/test_flows.py
index c92224ce..9e3d888f 100644
--- a/tests/test_flows.py
+++ b/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

From d391b3244dc5c6afa28c02e1c6d497ffa378e553 Mon Sep 17 00:00:00 2001
From: Noah Wolfe <newolfe@mit.edu>
Date: Mon, 9 Feb 2026 16:30:03 -0500
Subject: [PATCH 08/10] add lineax as dependency for efficient numerical
 inverse jvp

---
 pyproject.toml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pyproject.toml b/pyproject.toml
index 06d1291a..d9cd03ef 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -16,6 +16,7 @@ dependencies = [
     "tqdm",
     "optax",
     "paramax",
+    "lineax"
 ]
 description = "Easy to use distributions, bijections and normalizing flows in JAX."
 keywords = ["jax", "neural-networks", "equinox"]

From 9f67a9746c9f9a61ceb8ac7e8413e0612668f153 Mon Sep 17 00:00:00 2001
From: Noah Wolfe <newolfe@mit.edu>
Date: Tue, 17 Feb 2026 11:26:20 -0500
Subject: [PATCH 09/10] simplify NumericalInverse implicit diff behavior with
 use_implicit_differentation flag; remove unused make_layer in
 BlockAutoregressiveNetwork __init__

---
 .../block_autoregressive_network.py           |  4 ----
 flowjax/bijections/utils.py                   | 19 +++++++------------
 2 files changed, 7 insertions(+), 16 deletions(-)

diff --git a/flowjax/bijections/block_autoregressive_network.py b/flowjax/bijections/block_autoregressive_network.py
index e1484de5..8a15c8e5 100644
--- a/flowjax/bijections/block_autoregressive_network.py
+++ b/flowjax/bijections/block_autoregressive_network.py
@@ -122,10 +122,6 @@ def __init__(
                 (1, block_dim),
             ]
 
-            def make_layer(inp):
-                layer_key, block_shape = inp
-                return block_autoregressive_linear(layer_key, n_blocks=dim, block_shape=block_shape)
-
             for layer_key, block_shape in zip(keys, block_shapes, strict=True):
                 layers_and_log_jac_fns.append(
                     block_autoregressive_linear(
diff --git a/flowjax/bijections/utils.py b/flowjax/bijections/utils.py
index bcd8bf62..6721ced2 100644
--- a/flowjax/bijections/utils.py
+++ b/flowjax/bijections/utils.py
@@ -287,12 +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.
-        diffable_inverter: If ``False`` (default), gradients through the numerical
+        use_implicit_differentation: If ``True`` (default), gradients through the numerical
             inverse are supplied with a custom JVP using implicit differentiation. If
-            ``True``, the inverter is assumed to be safely differentiable as provided.
-        raise_old_error: If ``True``, preserve legacy behaviour and raise
-            ``RuntimeError`` when users attempt to differentiate through the numerical
-            inverse. Takes precedence over ``diffable_inverter``.
+            ``False``, the inverter is assumed to be safely differentiable as provided.
     """
 
     bijection: AbstractBijection
@@ -304,18 +301,16 @@ def __init__(
         self,
         bijection: AbstractBijection,
         inverter: Callable[[AbstractBijection, Array, Array | None], Array],
-        diffable_inverter: bool = False,
-        raise_old_error: bool = False,
+        use_implicit_differentation: bool = True,
     ):
         self.bijection = bijection
         self.shape = self.bijection.shape
         self.cond_shape = self.bijection.cond_shape
-        if raise_old_error:
-            self.inverter = self._wrap_inverter_with_error_on_grad(inverter)
-        elif diffable_inverter:
-            self.inverter = inverter
-        else:
+
+        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):

From 6b1def2e271985d372f9fc18d30b0ae27243b652 Mon Sep 17 00:00:00 2001
From: Noah Wolfe <newolfe@mit.edu>
Date: Tue, 17 Feb 2026 12:31:11 -0500
Subject: [PATCH 10/10] pass through use_implicit_differentation flag and
 remove test case for `raise_old_error` flag

---
 flowjax/flows.py                              | 15 +++++----------
 tests/test_bijections/test_bijection_utils.py | 16 ----------------
 2 files changed, 5 insertions(+), 26 deletions(-)

diff --git a/flowjax/flows.py b/flowjax/flows.py
index cc01658b..70394d9f 100644
--- a/flowjax/flows.py
+++ b/flowjax/flows.py
@@ -168,8 +168,7 @@ def block_neural_autoregressive_flow(
     invert: bool = True,
     activation: AbstractBijection | Callable | None = None,
     inverter: Callable[[AbstractBijection, Array, Array | None], Array] | None = None,
-    diffable_inverter: bool = False,
-    raise_old_error: bool = False
+    use_implicit_differentation: bool = True
 ) -> Transformed:
     """Block neural autoregressive flow (BNAF) (https://arxiv.org/abs/1904.04676).
 
@@ -200,12 +199,9 @@ def block_neural_autoregressive_flow(
             invert the ``BlockAutoregressiveNetwork`` bijection. Passed to
             :py:class:`~flowjax.bijections.NumericalInverse`. Defaults to
             using ``elementwise_autoregressive_bisection``.
-        diffable_inverter: Passed to :py:class:`~flowjax.bijections.NumericalInverse`.
-            If ``False`` (default), gradients through the inverse use a custom JVP.
-            If ``True``, the provided ``inverter`` is treated as differentiable.
-        raise_old_error: Passed to :py:class:`~flowjax.bijections.NumericalInverse`.
-            If ``True``, attempting to differentiate through the numerical inverse
-            raises a ``RuntimeError`` (legacy behavior).
+        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]
 
@@ -230,8 +226,7 @@ def make_layer(key):  # bnaf layer + permutation
                 activation=activation,
             ),
             inverter=inverter,
-            diffable_inverter=diffable_inverter,
-            raise_old_error=raise_old_error
+            use_implicit_differentation=use_implicit_differentation
         )
         return _add_default_permute(bijection, base_dist.shape[-1], perm_key)
 
diff --git a/tests/test_bijections/test_bijection_utils.py b/tests/test_bijections/test_bijection_utils.py
index 22316f80..007c47c9 100644
--- a/tests/test_bijections/test_bijection_utils.py
+++ b/tests/test_bijections/test_bijection_utils.py
@@ -51,19 +51,3 @@ def test_NumericalInverse_implicit_jvp(diff_fn):
     )
     derivative = diff_fn(bijection.inverse)(jnp.array(1.0))
     assert jnp.allclose(derivative, 0.5, atol=1e-5)
-
-
-@pytest.mark.parametrize("diff_fn", test_cases)
-def test_NumericalInverse_raise_old_error_takes_precedence(diff_fn):
-    bijection = NumericalInverse(
-        Affine(5, 2),
-        root_finder_to_inverter(
-            partial(bisection_search, lower=-1, upper=1, atol=1e-7),
-        ),
-        diffable_inverter=True,
-        raise_old_error=True,
-    )
-    with pytest.raises(
-        RuntimeError, match="Computing gradients through the numerical inverse"
-    ):
-        diff_fn(bijection.inverse)(jnp.ones(()))