diff --git a/flowjax/bijections/rational_quadratic_spline.py b/flowjax/bijections/rational_quadratic_spline.py index 41af1ab4..071ab957 100644 --- a/flowjax/bijections/rational_quadratic_spline.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index b92f73a2..4bf79db7 100644 --- a/pyproject.toml +++ b/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" diff --git a/tests/test_bijections/test_rational_quadratic_spline.py b/tests/test_bijections/test_rational_quadratic_spline.py index f203aa85..cdf3d68a 100644 --- a/tests/test_bijections/test_rational_quadratic_spline.py +++ b/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"