From c68ade899a68d88000be7487d8682b38945dfc5e Mon Sep 17 00:00:00 2001 From: Christoph Schoenle Date: Thu, 18 Jun 2026 09:12:42 +0000 Subject: [PATCH 1/2] Fix: resolve NaN gradients in RQS out-of-bounds inverse pass --- .../bijections/rational_quadratic_spline.py | 12 +++-- .../test_rational_quadratic_spline.py | 46 ++++++++++++++++++- 2 files changed, 52 insertions(+), 6 deletions(-) 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/tests/test_bijections/test_rational_quadratic_spline.py b/tests/test_bijections/test_rational_quadratic_spline.py index f203aa85..e8f56a6a 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_map, tree_leaves from flowjax.bijections import RationalQuadraticSpline @@ -36,3 +36,47 @@ def test_RationalQuadraticSpline_init(interval): spline = RationalQuadraticSpline(knots=10, interval=interval) y = vmap(spline.transform)(x) assert pytest.approx(x, abs=1e-6) == y + +@pytest.mark.parametrize("direction", ["forward", "inverse"]) +def test_RationalQuadraticSpline_out_of_bounds_grad_nan(direction): + """Checks that gradients do not contain NaNs when evaluating the forward + and inverse passes on out-of-bounds inputs.""" + + # 1. Initialize RQS + spline = RationalQuadraticSpline(knots=4, interval=1.0) + + # 2. Partition and add truly random noise to simulate a "trained" state + params, static = eqx.partition(spline, eqx.is_inexact_array) + key = jr.key(42) + + def add_noise(x): + nonlocal key + if x is not None: + key, subkey = jr.split(key) + return x + jr.normal(subkey, x.shape) * 1.5 + return x + + params_noisy = tree_map(add_noise, params) + spline_trained = eqx.combine(params_noisy, static) + + # 3. Provide an out-of-bounds input (interval is 1.0, so 50.0 is way out) + val = jnp.array(50.0) + + # 4. Define a dummy loss evaluating the requested direction + def loss(model, v): + if direction == "forward": + out, _ = model.transform_and_log_det(v) + else: + out, _ = model.inverse_and_log_det(v) + return out.sum() + + # 5. Compute gradients + _, grad = eqx.filter_value_and_grad(loss)(spline_trained, val) + + # 6. Check for NaNs + nan_found = any( + jnp.any(jnp.isnan(leaf)) + for leaf in tree_leaves(grad) if eqx.is_array(leaf) + ) + + assert not nan_found, f"NaNs detected in gradients for out-of-bounds {direction} pass" \ No newline at end of file From ac3178dfde1da672bcbf6936aaa99e9ff9e545be Mon Sep 17 00:00:00 2001 From: Daniel Ward Date: Sun, 21 Jun 2026 19:14:22 +0100 Subject: [PATCH 2/2] Add version bump and simplify test --- pyproject.toml | 2 +- .../test_rational_quadratic_spline.py | 62 ++++++------------- 2 files changed, 21 insertions(+), 43 deletions(-) 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 e8f56a6a..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, tree_leaves +from jax.tree_util import tree_leaves, tree_map from flowjax.bijections import RationalQuadraticSpline @@ -37,46 +37,24 @@ def test_RationalQuadraticSpline_init(interval): y = vmap(spline.transform)(x) assert pytest.approx(x, abs=1e-6) == y -@pytest.mark.parametrize("direction", ["forward", "inverse"]) -def test_RationalQuadraticSpline_out_of_bounds_grad_nan(direction): - """Checks that gradients do not contain NaNs when evaluating the forward - and inverse passes on out-of-bounds inputs.""" - - # 1. Initialize RQS - spline = RationalQuadraticSpline(knots=4, interval=1.0) - - # 2. Partition and add truly random noise to simulate a "trained" state - params, static = eqx.partition(spline, eqx.is_inexact_array) - key = jr.key(42) - - def add_noise(x): - nonlocal key - if x is not None: - key, subkey = jr.split(key) - return x + jr.normal(subkey, x.shape) * 1.5 - return x - - params_noisy = tree_map(add_noise, params) - spline_trained = eqx.combine(params_noisy, static) - - # 3. Provide an out-of-bounds input (interval is 1.0, so 50.0 is way out) - val = jnp.array(50.0) - - # 4. Define a dummy loss evaluating the requested direction - def loss(model, v): - if direction == "forward": - out, _ = model.transform_and_log_det(v) - else: - out, _ = model.inverse_and_log_det(v) - return out.sum() - - # 5. Compute gradients - _, grad = eqx.filter_value_and_grad(loss)(spline_trained, val) - # 6. Check for NaNs - nan_found = any( - jnp.any(jnp.isnan(leaf)) - for leaf in tree_leaves(grad) if eqx.is_array(leaf) +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 ) - - assert not nan_found, f"NaNs detected in gradients for out-of-bounds {direction} pass" \ No newline at end of file + 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"