Fix: resolve NaN gradients in RQS out-of-bounds inverse pass

flowjax/bijections/rational_quadratic_spline.py

@@ -66,9 +66,10 @@ def transform_and_log_det(self, x, condition=None):
             x >= jnp.array(self.interval[0]),
             x <= jnp.array(self.interval[1]),
         )
-        k = jnp.searchsorted(x_pos, x) - 1  # k is bin number
+        x_clipped = jnp.clip(x, self.interval[0], self.interval[1])
+        k = jnp.searchsorted(x_pos, x_clipped) - 1  # k is bin number
         k = jnp.clip(k, min=0, max=len(x_pos) - 2)
-        xi = (x - x_pos[k]) / (x_pos[k + 1] - x_pos[k])
+        xi = (x_clipped - x_pos[k]) / (x_pos[k + 1] - x_pos[k])
         sk = (y_pos[k + 1] - y_pos[k]) / (x_pos[k + 1] - x_pos[k])
         dk, dk1, yk, yk1 = derivatives[k], derivatives[k + 1], y_pos[k], y_pos[k + 1]
         num = (yk1 - yk) * (sk * xi**2 + dk * xi * (1 - xi))
@@ -94,14 +95,15 @@ def inverse_and_log_det(self, y, condition=None):
             y >= jnp.array(self.interval[0]),
             y <= jnp.array(self.interval[1]),
         )
-        k = jnp.searchsorted(y_pos, y) - 1
+        y_clipped = jnp.clip(y, self.interval[0], self.interval[1])
+        k = jnp.searchsorted(y_pos, y_clipped) - 1
         k = jnp.clip(k, min=0, max=len(x_pos) - 2)
         xk, xk1, yk, yk1 = x_pos[k], x_pos[k + 1], y_pos[k], y_pos[k + 1]
         sk = (yk1 - yk) / (xk1 - xk)
-        y_delta_s_term = (y - yk) * (derivatives[k + 1] + derivatives[k] - 2 * sk)
+        y_delta_s_term = (y_clipped - yk) * (derivatives[k + 1] + derivatives[k] - 2 * sk)
         a = (yk1 - yk) * (sk - derivatives[k]) + y_delta_s_term
         b = (yk1 - yk) * derivatives[k] - y_delta_s_term
-        c = -sk * (y - yk)
+        c = -sk * (y_clipped - yk)
         sqrt_term = jnp.sqrt(b**2 - 4 * a * c)
         xi = (2 * c) / (-b - sqrt_term)
         x = xi * (xk1 - xk) + xk

pyproject.toml

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

tests/test_bijections/test_rational_quadratic_spline.py

@@ -3,7 +3,7 @@
 import jax.random as jr
 import pytest
 from jax import vmap
-from jax.tree_util import tree_map
+from jax.tree_util import tree_leaves, tree_map
 
 from flowjax.bijections import RationalQuadraticSpline
 
@@ -36,3 +36,25 @@ def test_RationalQuadraticSpline_init(interval):
     spline = RationalQuadraticSpline(knots=10, interval=interval)
     y = vmap(spline.transform)(x)
     assert pytest.approx(x, abs=1e-6) == y
+
+
+def test_RationalQuadraticSpline_out_of_bounds_grad_nan():
+    """Checks that gradients do not contain NaNs for out of interval inputs."""
+    key = jr.key(1)
+    spline = RationalQuadraticSpline(knots=4, interval=1.0)
+    spline = tree_map(  # Not identity
+        lambda x: x + jr.normal(key, x.shape) * 1.5 if eqx.is_inexact_array(x) else x,
+        spline,
+    )
+    val = jnp.array(50.0)  # Out of bounds
+    forward_grad = eqx.filter_grad(lambda m, v: m.transform_and_log_det(v)[0].sum())(
+        spline, val
+    )
+    inverse_grad = eqx.filter_grad(lambda m, v: m.inverse_and_log_det(v)[0].sum())(
+        spline, val
+    )
+
+    for grad in (forward_grad, inverse_grad):
+        assert not any(
+            jnp.any(jnp.isnan(leaf)) for leaf in tree_leaves(grad) if eqx.is_array(leaf)
+        ), "NaNs detected in gradients for out-of-bounds pass"