From 98b5be051f126c8e9d3e1949bf273d7b6fe026b8 Mon Sep 17 00:00:00 2001 From: Gary Allen Date: Wed, 1 Jul 2026 22:57:11 +0200 Subject: [PATCH 1/2] Add Exp and Softplus bijectors Exponential and softplus bijectors modeled after TensorFlow Probability's implementations, mapping the real line to the positive domain. --- distreqx/bijectors/__init__.py | 2 + distreqx/bijectors/_exp.py | 34 +++++++++++++++++ distreqx/bijectors/_softplus.py | 40 ++++++++++++++++++++ docs/api/bijectors/exp.md | 6 +++ docs/api/bijectors/softplus.md | 6 +++ mkdocs.yml | 2 + tests/conftest.py | 5 +++ tests/exp_test.py | 65 +++++++++++++++++++++++++++++++++ tests/softplus_test.py | 36 ++++++++++++++++++ 9 files changed, 196 insertions(+) create mode 100644 distreqx/bijectors/_exp.py create mode 100644 distreqx/bijectors/_softplus.py create mode 100644 docs/api/bijectors/exp.md create mode 100644 docs/api/bijectors/softplus.md create mode 100644 tests/conftest.py create mode 100644 tests/exp_test.py create mode 100644 tests/softplus_test.py diff --git a/distreqx/bijectors/__init__.py b/distreqx/bijectors/__init__.py index 7805142..4054157 100644 --- a/distreqx/bijectors/__init__.py +++ b/distreqx/bijectors/__init__.py @@ -7,10 +7,12 @@ from ._block import Block as Block from ._chain import Chain as Chain from ._diag_linear import DiagLinear as DiagLinear +from ._exp import Exp as Exp from ._linear import AbstractLinearBijector as AbstractLinearBijector from ._scalar_affine import ScalarAffine as ScalarAffine from ._shift import Shift as Shift from ._sigmoid import Sigmoid as Sigmoid +from ._softplus import Softplus as Softplus from ._tanh import Tanh as Tanh from ._triangular_linear import TriangularLinear as TriangularLinear from ._unconstrained_affine import UnconstrainedAffine as UnconstrainedAffine diff --git a/distreqx/bijectors/_exp.py b/distreqx/bijectors/_exp.py new file mode 100644 index 0000000..7c8ed53 --- /dev/null +++ b/distreqx/bijectors/_exp.py @@ -0,0 +1,34 @@ +import jax.numpy as jnp +from jaxtyping import Array + +from ._bijector import ( + AbstractBijector, + AbstractForwardInverseBijector, + AbstractFwdLogDetJacBijector, + AbstractInvLogDetJacBijector, +) + + +class Exp( + AbstractForwardInverseBijector, + AbstractInvLogDetJacBijector, + AbstractFwdLogDetJacBijector, +): + """Exponential bijector: y = exp(x).""" + + _is_constant_jacobian: bool = False + _is_constant_log_det: bool = False + + def forward_and_log_det(self, x: Array) -> tuple[Array, Array]: + """Computes y = exp(x) and log|det J(f)(x)| = x.""" + return jnp.exp(x), x + + def inverse_and_log_det(self, y: Array) -> tuple[Array, Array]: + """Computes x = log(y) and log|det J(f^{-1})(y)| = -log(y).""" + x = jnp.log(y) + # Optimization: since x = log(y), the log det is simply -x + return x, -x + + def same_as(self, other: AbstractBijector) -> bool: + """Returns True if this bijector is guaranteed to be the same as `other`.""" + return type(other) is Exp diff --git a/distreqx/bijectors/_softplus.py b/distreqx/bijectors/_softplus.py new file mode 100644 index 0000000..55a87dd --- /dev/null +++ b/distreqx/bijectors/_softplus.py @@ -0,0 +1,40 @@ +import jax.nn as jnn +import jax.numpy as jnp +from jaxtyping import Array + +from ._bijector import ( + AbstractBijector, + AbstractForwardInverseBijector, + AbstractFwdLogDetJacBijector, + AbstractInvLogDetJacBijector, +) + + +class Softplus( + AbstractForwardInverseBijector, + AbstractInvLogDetJacBijector, + AbstractFwdLogDetJacBijector, +): + """ + Transforms the real line to the positive domain using + softplus y = log(1 + exp(x)). + """ + + _is_constant_jacobian: bool = False + _is_constant_log_det: bool = False + + def forward_and_log_det(self, x: Array) -> tuple[Array, Array]: + """Computes y = softplus(x) and log|det J(f)(x)|.""" + y = jnn.softplus(x) + logdet = -jnn.softplus(-x) + return y, logdet + + def inverse_and_log_det(self, y: Array) -> tuple[Array, Array]: + """Computes x = softplus^{-1}(y) and log|det J(f^{-1})(y)|.""" + x = jnp.log(-jnp.expm1(-y)) + y + logdet = jnn.softplus(-x) + return x, logdet + + def same_as(self, other: AbstractBijector) -> bool: + """Returns True if this bijector is guaranteed to be the same as `other`.""" + return type(other) is Softplus diff --git a/docs/api/bijectors/exp.md b/docs/api/bijectors/exp.md new file mode 100644 index 0000000..d103c6b --- /dev/null +++ b/docs/api/bijectors/exp.md @@ -0,0 +1,6 @@ +# Exp Bijector + +::: distreqx.bijectors.Exp + options: + members: false +--- diff --git a/docs/api/bijectors/softplus.md b/docs/api/bijectors/softplus.md new file mode 100644 index 0000000..0087d64 --- /dev/null +++ b/docs/api/bijectors/softplus.md @@ -0,0 +1,6 @@ +# Softplus Bijector + +::: distreqx.bijectors.Softplus + options: + members: false +--- diff --git a/mkdocs.yml b/mkdocs.yml index aa5cf4e..d0fe279 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -124,6 +124,8 @@ nav: - 'api/bijectors/shift.md' - 'api/bijectors/sigmoid.md' - 'api/bijectors/tanh.md' + - 'api/bijectors/exp.md' + - 'api/bijectors/softplus.md' - 'api/bijectors/triangular_linear.md' - 'api/bijectors/_bijector.md' - Utilities: diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..8f407ab --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,5 @@ +import jax + +# Must be set before any JAX arrays are initialized. Living here (rather than in +# individual test files) guarantees it runs before any test module import. +jax.config.update("jax_enable_x64", True) diff --git a/tests/exp_test.py b/tests/exp_test.py new file mode 100644 index 0000000..47fd06a --- /dev/null +++ b/tests/exp_test.py @@ -0,0 +1,65 @@ +from unittest import TestCase + +import equinox as eqx +import jax +import jax.numpy as jnp +import numpy as np +from parameterized import parameterized # type: ignore + +from distreqx.bijectors import Exp, Tanh + + +class ExpTest(TestCase): + def setUp(self): + self.bij = Exp() + + def assertion_fn(self, rtol=1e-5): + return lambda x, y: np.testing.assert_allclose(x, y, rtol=rtol) + + @parameterized.expand([("float32", jnp.float32), ("float64", jnp.float64)]) + def test_forward_and_log_det(self, name, dtype): + x = jnp.array([-2.5, 0.0, 1.0, 3.5], dtype=dtype) + y, log_det = self.bij.forward_and_log_det(x) + + self.assertion_fn()(y, jnp.exp(x)) + self.assertion_fn()(log_det, x) + self.assertEqual(y.dtype, dtype) + self.assertEqual(log_det.dtype, dtype) + + @parameterized.expand([("float32", jnp.float32), ("float64", jnp.float64)]) + def test_inverse_and_log_det(self, name, dtype): + # We must use strictly positive numbers for the domain of log(y) + y = jnp.array([0.1, 1.0, jnp.e, 10.0], dtype=dtype) + x, log_det = self.bij.inverse_and_log_det(y) + + self.assertion_fn()(x, jnp.log(y)) + self.assertion_fn()(log_det, -jnp.log(y)) + self.assertEqual(x.dtype, dtype) + self.assertEqual(log_det.dtype, dtype) + + @parameterized.expand([("float32", jnp.float32), ("float64", jnp.float64)]) + def test_forward_and_inverse(self, name, dtype): + # Round-trip: inverse(forward(x)) == x, matching the softplus tests. + x = jnp.array([-5.0, 0.0, 5.0], dtype=dtype) + y, log_det_fwd = self.bij.forward_and_log_det(x) + + x_rec, log_det_inv = self.bij.inverse_and_log_det(y) + self.assertion_fn()(x_rec, x) + self.assertion_fn()(log_det_inv, -log_det_fwd) + + def test_jittable(self): + @eqx.filter_jit + def f(bij, x): + return bij.forward_and_log_det(x) + + x = jnp.array([1.0, 2.0]) + y, log_det = f(self.bij, x) + self.assertIsInstance(y, jax.Array) + self.assertIsInstance(log_det, jax.Array) + + def test_same_as(self): + same_bij = Exp() + diff_bij = Tanh() + + self.assertTrue(self.bij.same_as(same_bij)) + self.assertFalse(self.bij.same_as(diff_bij)) diff --git a/tests/softplus_test.py b/tests/softplus_test.py new file mode 100644 index 0000000..c13036d --- /dev/null +++ b/tests/softplus_test.py @@ -0,0 +1,36 @@ +from unittest import TestCase + +import equinox as eqx +import jax +import jax.numpy as jnp +import numpy as np + +from distreqx.bijectors import Softplus + + +class SoftplusTest(TestCase): + def setUp(self): + self.bij = Softplus() + + def assertion_fn(self, rtol=1e-5): + return lambda x, y: np.testing.assert_allclose(x, y, rtol=rtol) + + def test_forward_and_inverse(self): + x = jnp.array([-5.0, 0.0, 5.0]) + y, log_det_fwd = self.bij.forward_and_log_det(x) + + expected_y = jax.nn.softplus(x) + self.assertion_fn()(y, expected_y) + self.assertion_fn()(log_det_fwd, -jax.nn.softplus(-x)) + + x_rec, log_det_inv = self.bij.inverse_and_log_det(y) + self.assertion_fn()(x_rec, x) + self.assertion_fn()(log_det_inv, -log_det_fwd) + + def test_jittable(self): + @eqx.filter_jit + def f(bij, x): + return bij.forward_and_log_det(x) + + y, log_det = f(self.bij, jnp.array(1.0)) + self.assertIsInstance(y, jax.Array) From 3fb633b4ff724511aa4914eb10a7b6d8a2bf816a Mon Sep 17 00:00:00 2001 From: Gary Allen Date: Wed, 8 Jul 2026 21:14:51 +0200 Subject: [PATCH 2/2] Fix NaN gradient in Sigmoid's stable sigmoid/softplus helpers _more_stable_sigmoid and _more_stable_softplus used jnp.where(cond, a, b) to switch to a numerically stable approximation for very negative inputs. jnp.where evaluates both branches unconditionally, so for large positive inputs the unselected branch (jnp.exp(x) / log1p(exp(x))) overflowed to inf, and the 0 * inf this produces in the backward pass poisoned the gradient with NaN even though that branch was never used for the forward value. Fix by feeding the unselected branch a safe input via a second jnp.where, the standard remedy for this "double where" trap. Forward values are unchanged; only the previously-poisoned gradient at large positive inputs is fixed. Adds a regression test exercising both helpers (via Sigmoid.forward and Sigmoid.forward_log_det_jacobian) across a range including extreme positive and negative inputs. --- distreqx/bijectors/_sigmoid.py | 13 +++++++++++-- tests/sigmoid_test.py | 18 ++++++++++++++++++ 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/distreqx/bijectors/_sigmoid.py b/distreqx/bijectors/_sigmoid.py index f38b93b..73d1ee9 100644 --- a/distreqx/bijectors/_sigmoid.py +++ b/distreqx/bijectors/_sigmoid.py @@ -63,7 +63,13 @@ def same_as(self, other: AbstractBijector) -> bool: def _more_stable_sigmoid(x: Array) -> Array: """Where extremely negatively saturated, approximate sigmoid with exp(x).""" - ret = jnp.where(x < -9, jnp.exp(x), jax.nn.sigmoid(x)) + # `jnp.where` evaluates both branches, so the unselected `jnp.exp(x)` branch + # would overflow to +inf for large positive `x`, and the resulting `0 * inf` + # in the backward pass poisons the gradient with NaN (the classic JAX + # "double where" trap). Feeding the unselected branch a safe input keeps the + # forward value identical while making its gradient finite. + safe_x = jnp.where(x < -9, x, 0.0) + ret = jnp.where(x < -9, jnp.exp(safe_x), jax.nn.sigmoid(x)) if not isinstance(ret, Array): raise TypeError("ret is not an Array") return ret @@ -71,7 +77,10 @@ def _more_stable_sigmoid(x: Array) -> Array: def _more_stable_softplus(x: Array) -> Array: """Where extremely saturated, approximate softplus with log1p(exp(x)).""" - ret = jnp.where(x < -9, jnp.log1p(jnp.exp(x)), jax.nn.softplus(x)) + # See `_more_stable_sigmoid`: guard the unselected `log1p(exp(x))` branch + # against a +inf-poisoned NaN gradient for large positive `x`. + safe_x = jnp.where(x < -9, x, 0.0) + ret = jnp.where(x < -9, jnp.log1p(jnp.exp(safe_x)), jax.nn.softplus(x)) if not isinstance(ret, Array): raise TypeError("ret is not an Array") return ret diff --git a/tests/sigmoid_test.py b/tests/sigmoid_test.py index 5aebf14..21eb98a 100644 --- a/tests/sigmoid_test.py +++ b/tests/sigmoid_test.py @@ -101,6 +101,24 @@ def test_inverse_and_log_det(self): np.testing.assert_allclose(x1, x2, rtol=RTOL) np.testing.assert_allclose(logdet1, logdet2, rtol=RTOL) + def test_gradient_finite_at_extreme_values(self): + """Regression test for a NaN-gradient bug in `_more_stable_sigmoid`/ + `_more_stable_softplus`. Both used a single `jnp.where(cond, a, b)` to + pick between two approximations, but `jnp.where` evaluates both branches + unconditionally: for large positive `x`, the unselected `jnp.exp(x)` (or + `log1p(exp(x))`) branch overflowed to `inf`, and the `0 * inf` this + produces in the backward pass poisoned the gradient with NaN even though + that branch was never selected for the forward value. + """ + bijector = Sigmoid() + x = jnp.array([-1e5, -50.0, -9.0, 0.0, 9.0, 50.0, 1e5]) + + fwd_grad = jax.vmap(jax.grad(bijector.forward))(x) + self.assertTrue(bool(jnp.all(jnp.isfinite(fwd_grad)))) + + logdet_grad = jax.vmap(jax.grad(bijector.forward_log_det_jacobian))(x) + self.assertTrue(bool(jnp.all(jnp.isfinite(logdet_grad)))) + def test_jittable(self): @jax.jit def f(x, b):