Add Exp and Softplus

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

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

distreqx/bijectors/_sigmoid.py

@@ -63,15 +63,24 @@ 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
 
 
 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

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

docs/api/bijectors/exp.md

@@ -0,0 +1,6 @@
+# Exp Bijector
+
+::: distreqx.bijectors.Exp
+    options:
+        members: false
+---

docs/api/bijectors/softplus.md

@@ -0,0 +1,6 @@
+# Softplus Bijector
+
+::: distreqx.bijectors.Softplus
+    options:
+        members: false
+---

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'

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))

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):

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)