diff --git a/flowjax/bijections/utils.py b/flowjax/bijections/utils.py index d906a0fe..6721ced2 100644 --- a/flowjax/bijections/utils.py +++ b/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 diff --git a/flowjax/flows.py b/flowjax/flows.py index 231b310e..70394d9f 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, + 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) diff --git a/flowjax/masks.py b/flowjax/masks.py index 71a1ced1..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,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) diff --git a/pyproject.toml b/pyproject.toml index d4bb290c..a11766b9 100644 --- a/pyproject.toml +++ b/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"] diff --git a/tests/test_bijections/test_bijection_utils.py b/tests/test_bijections/test_bijection_utils.py index f03540c8..007c47c9 100644 --- a/tests/test_bijections/test_bijection_utils.py +++ b/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) 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