Merge pull request #219 from danielward27/better_errors

Better errors

Author: danielward27

flowjax/bijections/utils.py

@@ -293,8 +293,23 @@ def __init__(
         bijection: AbstractBijection,
         inverter: Callable[[AbstractBijection, Array, Array | None], Array],
     ):
+        @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.bijection = bijection
-        self.inverter = inverter
+        self.inverter = nondiff_inverter
         self.shape = self.bijection.shape
         self.cond_shape = self.bijection.cond_shape
 

flowjax/distributions.py

@@ -244,6 +244,27 @@ class AbstractTransformed(AbstractDistribution):
     base_dist: AbstractVar[AbstractDistribution]
     bijection: AbstractVar[AbstractBijection]
 
+    def __check_init__(self):
+        """Check for compatible shapes between base_dist and bijection."""
+        if (
+            self.base_dist.cond_shape is not None
+            and self.bijection.cond_shape is not None
+            and self.base_dist.cond_shape != self.bijection.cond_shape
+        ):
+            raise ValueError(
+                "The base distribution and bijection are both conditional "
+                "but have mismatched cond_shape attributes. Base distribution has"
+                f"{self.base_dist.cond_shape}, and the bijection has"
+                f"{self.bijection.cond_shape}.",
+            )
+
+        if self.base_dist.shape != self.bijection.shape:
+            raise ValueError(
+                "The base distribution and bijection have mismatched shapes. "
+                f"Base distribution has {self.base_dist.shape}, and the bijection "
+                f"has {self.bijection.shape}.",
+            )
+
     def _log_prob(self, x, condition=None):
         z, log_abs_det = self.bijection.inverse_and_log_det(x, condition)
         p_z = self.base_dist._log_prob(z, condition)
@@ -268,20 +289,6 @@ def _sample_and_log_prob(
         )
         return sample, log_prob_base - forward_log_dets
 
-    def __check_init__(self):  # TODO test errors and test conditional base distribution
-        """Checks cond_shape is compatible in both bijection and distribution."""
-        if (
-            self.base_dist.cond_shape is not None
-            and self.bijection.cond_shape is not None
-            and self.base_dist.cond_shape != self.bijection.cond_shape
-        ):
-            raise ValueError(
-                "The base distribution and bijection are both conditional "
-                "but have mismatched cond_shape attributes. Base distribution has"
-                f"{self.base_dist.cond_shape}, and the bijection has"
-                f"{self.bijection.cond_shape}.",
-            )
-
     def merge_transforms(self):
         """Unnests nested transformed distributions.
 

pyproject.toml

@@ -23,7 +23,7 @@ license = { file = "LICENSE" }
 name = "flowjax"
 readme = "README.md"
 requires-python = ">=3.10"
-version = "17.1.1"
+version = "17.1.2"
 
 [project.urls]
 repository = "https://github.com/danielward27/flowjax"

tests/test_bijections/test_bijection_utils.py

@@ -1,8 +1,12 @@
+from functools import partial
+
+import jax
 import jax.numpy as jnp
 import pytest
 from equinox import EquinoxRuntimeError
 
-from flowjax.bijections import Affine, Indexed, Permute
+from flowjax.bijections import Affine, Indexed, NumericalInverse, Permute
+from flowjax.root_finding import bisection_search, root_finder_to_inverter
 
 test_cases = {
     # name: idx, expected
@@ -32,3 +36,20 @@ def test_partial(idx, expected):
 def test_Permute_argcheck():
     with pytest.raises(EquinoxRuntimeError):
         Permute(jnp.array([0, 0]))
+
+
+test_cases = [jax.grad, jax.jacfwd, jax.jacrev]
+
+
+@pytest.mark.parametrize("diff_fn", test_cases)
+def test_NumericalInverse_not_differentiable(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(()))

tests/test_distributions.py

@@ -120,6 +120,7 @@ def test_uniform_params():
 
 class _TestDist(AbstractDistribution):
     "Toy distribution object, for testing of distribution broadcasting."
+
     shape: tuple[int, ...]
     cond_shape: tuple[int, ...] | None
 
@@ -292,3 +293,12 @@ def test_transformed():
     assert dist.sample(jr.key(0), condition=jnp.ones((5, 2))).shape == (5,)
     assert dist.shape == ()
     assert dist.cond_shape == (2,)
+
+
+def test_transformed_wrong_shape():
+
+    with pytest.raises(ValueError, match="mismatched shapes"):
+        Transformed(
+            StandardNormal((2,)),
+            Affine(),
+        )