From 04802b9077f37cfe82e0c0d60d6ed1455cefb675 Mon Sep 17 00:00:00 2001 From: Sumukh Chaluvaraju Date: Thu, 4 Jun 2026 17:45:51 +0100 Subject: [PATCH 1/2] fix(tanh): clip inverse to (-1+eps, 1-eps) to prevent NaN in log_prob MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit In float32, sampling from a Tanh-transformed distribution (e.g. SAC or PPO with a squashed Gaussian policy) can produce values numerically equal to ±1 due to limited precision. Passing these to arctanh gives ±∞, which propagates to NaN in log_prob and silently breaks policy gradient updates. Fix: clip y to the open interval (-1+ε, 1-ε) in inverse_and_log_det, where ε = jnp.finfo(y.dtype).eps — the tightest safe bound for the dtype. This produces the largest finite pre-activation rather than ±∞ and keeps log_prob finite for any sample produced by sample(). The clip is dtype-aware so float64 benefits from the tighter 2.2e-16 bound while float32 uses 1.2e-7. Also updates test_stability to verify finiteness at boundary values (distrax intentionally diverges from TFP's NaN behaviour here) and adds two regression tests: one for direct boundary clipping, one that samples from a wide Tanh-wrapped MultivariateNormalDiag and verifies log_prob remains finite. Fixes: google-deepmind/distrax#216 --- distrax/_src/bijectors/tanh.py | 13 ++++++++- distrax/_src/bijectors/tanh_test.py | 45 +++++++++++++++++++++++++++-- 2 files changed, 55 insertions(+), 3 deletions(-) diff --git a/distrax/_src/bijectors/tanh.py b/distrax/_src/bijectors/tanh.py index 3b81dcb8..81422d9e 100644 --- a/distrax/_src/bijectors/tanh.py +++ b/distrax/_src/bijectors/tanh.py @@ -59,7 +59,18 @@ def forward_and_log_det(self, x: Array) -> Tuple[Array, Array]: return jnp.tanh(x), self.forward_log_det_jacobian(x) def inverse_and_log_det(self, y: Array) -> Tuple[Array, Array]: - """Computes x = f^{-1}(y) and log|det J(f^{-1})(y)|.""" + """Computes x = f^{-1}(y) and log|det J(f^{-1})(y)|. + + When `float32` is used, sampling from a Tanh-transformed distribution can + produce values numerically equal to ±1 due to limited precision, making + `arctanh` undefined and causing NaN in `log_prob`. We clip `y` to the + open interval `(-1 + ε, 1 - ε)` where `ε = jnp.finfo(y.dtype).eps`, which + is the tightest safe bound for the given dtype. This matches the workaround + documented in issue https://github.com/google-deepmind/distrax/issues/216 + and makes `log_prob` finite for any sample produced by `sample()`. + """ + eps = jnp.finfo(y.dtype).eps + y = jnp.clip(y, -1.0 + eps, 1.0 - eps) x = jnp.arctanh(y) # pyrefly: ignore[unsupported-operation] return x, -self.forward_log_det_jacobian(x) diff --git a/distrax/_src/bijectors/tanh_test.py b/distrax/_src/bijectors/tanh_test.py index 9810ef5e..f3995097 100644 --- a/distrax/_src/bijectors/tanh_test.py +++ b/distrax/_src/bijectors/tanh_test.py @@ -138,9 +138,16 @@ def test_stability(self): np.testing.assert_allclose(fldj_, fldj, rtol=RTOL) y = bijector.forward(x) # pytype: disable=wrong-arg-types # jax-ndarray - ildj = tfp_bijector.inverse_log_det_jacobian(y, event_ndims=0) + # For the inverse log-det, distrax clips boundary values to prevent NaN + # (unlike TFP which returns NaN for tanh(±10) = ±1 in float32). + # We verify finiteness rather than matching TFP's NaN for those entries. + # Interior values (|x| < 10) still agree with TFP. ildj_ = self.variant(bijector.inverse_log_det_jacobian)(y) # pyrefly: ignore[missing-attribute] - np.testing.assert_allclose(ildj_, ildj, rtol=RTOL) + self.assertFalse(np.any(np.isnan(ildj_)), + 'inverse_log_det_jacobian should be finite, got NaN') + interior = np.array([1, 2, 3], dtype=int) # indices for x in {-3.3,0,3.3} + ildj_tfp = tfp_bijector.inverse_log_det_jacobian(y, event_ndims=0) + np.testing.assert_allclose(ildj_[interior], ildj_tfp[interior], rtol=RTOL) @chex.all_variants @parameterized.named_parameters( @@ -173,6 +180,40 @@ def test_same_as(self): self.assertTrue(bijector.same_as(tanh.Tanh())) self.assertFalse(bijector.same_as(sigmoid.Sigmoid())) + def test_inverse_clips_boundary_values_to_prevent_nan(self): + """Regression test for https://github.com/google-deepmind/distrax/issues/216. + + In float32, sampling from a Tanh-transformed distribution can yield values + numerically equal to ±1 due to limited precision. arctanh(±1) = ±∞, + which causes NaN in log_prob. The fix clips y to (-1+eps, 1-eps). + """ + bijector = tanh.Tanh() + # Exact boundary values that would produce NaN without clipping. + y_boundary = jnp.array([-1.0, 1.0], dtype=jnp.float32) + x, log_det = bijector.inverse_and_log_det(y_boundary) + self.assertFalse(jnp.any(jnp.isnan(x)), 'x should be finite, got NaN') + self.assertFalse(jnp.any(jnp.isnan(log_det)), + 'log_det should be finite, got NaN') + self.assertFalse(jnp.any(jnp.isinf(x)), 'x should be finite, got Inf') + self.assertFalse(jnp.any(jnp.isinf(log_det)), + 'log_det should be finite, got Inf') + + def test_log_prob_finite_at_float32_boundary_samples(self): + """log_prob must be finite for samples that saturate float32 tanh.""" + import distrax + import jax.random as jr + # Build a Tanh-wrapped normal that is likely to produce boundary samples. + dist = distrax.Transformed( + distribution=distrax.MultivariateNormalDiag( + loc=jnp.zeros(4, dtype=jnp.float32), + scale_diag=jnp.ones(4, dtype=jnp.float32) * 10.0), # wide → ±1 + bijector=distrax.Block(distrax.Tanh(), ndims=1)) + key = jr.PRNGKey(0) + samples = dist.sample(seed=key, sample_shape=(16,)) + log_probs = dist.log_prob(samples) + self.assertFalse(jnp.any(jnp.isnan(log_probs)), + 'log_prob returned NaN for float32 boundary samples') + if __name__ == '__main__': jax.config.update('jax_threefry_partitionable', False) From e2223112a9a2e2d6b3784d4e41ed450043203a2d Mon Sep 17 00:00:00 2001 From: Sumukh Chaluvaraju Date: Thu, 4 Jun 2026 18:01:32 +0100 Subject: [PATCH 2/2] fix(mvn_from_bijector): make batch_shape robust to vmap-prepended dimensions When MultivariateNormal* distributions are constructed inside jax.vmap, the vmapped execution prepends a batch dimension to _loc at run time. The _batch_shape tuple is computed and captured at trace time (as a static Python tuple), so it does not include the extra batch dimension. This makes batch_shape stale and causes the loc property to fail with: ValueError: Cannot broadcast to shape with fewer dimensions: arr_shape=(B, D) shape=(D,) Fix: override batch_shape as a computed property. Instead of returning the static _batch_shape directly, detect any extra leading dimensions in _loc that go beyond what _batch_shape + _event_shape predict (i.e. dimensions added by vmap or similar batching transforms), and prepend them to the static batch_shape. This preserves the existing scale-broadcasting semantics (where batch_shape can be wider than loc.shape[:-1]) while also handling the vmap case. Also adds VmapBatchShapeTest with four regression cases: - batch_shape and loc.shape after vmap - loc property no longer raises after vmap (the reported symptom) - non-vmapped case unchanged - scale-batch-broadcasting unchanged Fixes: google-deepmind/distrax#276 --- .../_src/distributions/mvn_from_bijector.py | 28 +++++++++ .../distributions/mvn_from_bijector_test.py | 62 +++++++++++++++++++ 2 files changed, 90 insertions(+) diff --git a/distrax/_src/distributions/mvn_from_bijector.py b/distrax/_src/distributions/mvn_from_bijector.py index eb70a3ee..40fdc012 100644 --- a/distrax/_src/distributions/mvn_from_bijector.py +++ b/distrax/_src/distributions/mvn_from_bijector.py @@ -91,6 +91,34 @@ def __init__(self, loc: Array, scale: linear.Linear): self._batch_shape = batch_shape self._dtype = dtype + @property + def batch_shape(self): + """Returns the batch shape, correctly reflecting any extra leading dims. + + When the distribution is created inside ``jax.vmap`` (or any other JAX + transformation that prepends a batch dimension), the stored ``_loc`` + array acquires an extra leading dimension at run time that was not + present during tracing. The statically computed ``_batch_shape`` tuple + would then be too short, causing ``jnp.broadcast_to`` in ``loc`` to + fail with:: + + ValueError: Cannot broadcast to shape with fewer dimensions + + This property detects the extra dimensions by comparing the number of + leading dims in ``_loc`` against what ``_batch_shape`` and + ``_event_shape`` together predict, and prepends them if needed. + + See https://github.com/google-deepmind/distrax/issues/276. + """ + # Dimensions of _loc that should be accounted for by _batch_shape + _event_shape. + expected_ndim = len(self._batch_shape) + len(self._event_shape) + actual_ndim = len(self._loc.shape) + extra = actual_ndim - expected_ndim # >0 when vmap added batch dims + if extra > 0: + # Return the extra leading dims prepended to the static batch_shape. + return self._loc.shape[:extra] + self._batch_shape + return self._batch_shape + @property def scale(self) -> linear.Linear: """The scale bijector.""" diff --git a/distrax/_src/distributions/mvn_from_bijector_test.py b/distrax/_src/distributions/mvn_from_bijector_test.py index 0c92168d..7f26b061 100644 --- a/distrax/_src/distributions/mvn_from_bijector_test.py +++ b/distrax/_src/distributions/mvn_from_bijector_test.py @@ -280,6 +280,68 @@ def test_kl_divergence_raises_on_incompatible_distributions(self): dist1.kl_divergence(dist2) +class VmapBatchShapeTest(absltest.TestCase): + """Regression tests for https://github.com/google-deepmind/distrax/issues/276. + + When a MultivariateNormal* distribution is constructed inside jax.vmap, + the vmapped execution prepends a batch dimension to all JAX arrays, including + the stored `_loc`. The static `_batch_shape` tuple captured at trace time + does not include this extra dimension, causing `batch_shape` to be stale and + `loc` to raise: + ValueError: Cannot broadcast to shape with fewer dimensions + """ + + def test_batch_shape_after_vmap(self): + """batch_shape should reflect the vmap batch dimension.""" + @jax.jit + def build(): + def single(_): + return MultivariateNormalFromBijector( + loc=jnp.zeros(4), + scale=DiagLinear(diag=jnp.ones(4)), + ) + return jax.vmap(single)(jnp.arange(3)) + + dist = build() + self.assertEqual(dist.batch_shape, (3,)) + self.assertEqual(dist.event_shape, (4,)) + self.assertEqual(dist.loc.shape, (3, 4)) + + def test_loc_accessible_after_vmap(self): + """`loc` must not raise after vmap-construction (the reported symptom).""" + @jax.jit + def build(): + def single(_): + return MultivariateNormalFromBijector( + loc=jnp.zeros(4), + scale=DiagLinear(diag=jnp.ones(4)), + ) + return jax.vmap(single)(jnp.arange(5)) + + dist = build() + loc = dist.loc # must not raise + np.testing.assert_array_equal(loc, jnp.zeros((5, 4))) + + def test_non_vmapped_batch_shape_unchanged(self): + """Regression: static batch_shape must still be correct without vmap.""" + dist = MultivariateNormalFromBijector( + loc=jnp.zeros(4), + scale=DiagLinear(diag=jnp.ones(4)), + ) + self.assertEqual(dist.batch_shape, ()) + self.assertEqual(dist.loc.shape, (4,)) + + def test_scale_broadcasting_batch_shape_unchanged(self): + """Scale-batch broadcasting must still be respected without vmap.""" + # scale.batch_shape=(4,1), loc.batch=(1,3) → broadcast batch=(4,3) + dist = MultivariateNormalFromBijector( + loc=jnp.zeros((1, 3, 5)), + scale=DiagLinear(diag=jnp.ones((4, 1, 5))), + ) + self.assertEqual(dist.batch_shape, (4, 3)) + self.assertEqual(dist.loc.shape, (4, 3, 5)) + + if __name__ == '__main__': jax.config.update('jax_threefry_partitionable', False) absltest.main()