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