fix(distrax): Tanh NaN clip + MultivariateNormal vmap batch_shape fix

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)

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)

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."""

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