diff --git a/distreqx/bijectors/__init__.py b/distreqx/bijectors/__init__.py index 07e6e86..7b86059 100644 --- a/distreqx/bijectors/__init__.py +++ b/distreqx/bijectors/__init__.py @@ -7,12 +7,14 @@ 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 ._identity import Identity as Identity from ._inverse import Inverse as Inverse 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/_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/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 717446c..5730f19 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -122,10 +122,12 @@ nav: - 'api/bijectors/block.md' - 'api/bijectors/chain.md' - 'api/bijectors/diag_linear.md' + - 'api/bijectors/exp.md' - 'api/bijectors/scalar_affine.md' - 'api/bijectors/shift.md' - 'api/bijectors/sigmoid.md' - 'api/bijectors/tanh.md' + - 'api/bijectors/softplus.md' - 'api/bijectors/identity.md' - 'api/bijectors/inverse.md' - 'api/bijectors/triangular_linear.md' 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/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): 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)