Avoid intermediate NaN in multiply_no_nan.

distrax/_src/utils/math.py

@@ -43,7 +43,11 @@ def multiply_no_nan(x: Array, y: Array) -> Array:
     ValueError if the shapes of `x` and `y` do not match.
   """
   dtype = jnp.result_type(x, y)
-  return jnp.where(y == 0, jnp.zeros((), dtype=dtype), x * y)
+  # Replace `x` with zero where `y` is zero before multiplying, so that `0 * y`
+  # is computed instead of `x * 0`. This avoids producing an intermediate NaN
+  # when `x` is infinite and `y` is zero, which would be detected by
+  # `checkify`'s NaN checks even though the result is correct.
+  return jnp.where(y == 0, jnp.zeros((), dtype=dtype), x) * y
 
 
 # TODO(dougalm): move helpers like these into JAX AD utils

distrax/_src/utils/math_test.py

@@ -39,6 +39,18 @@ def test_multiply_no_nan_grads(self):
         lambda inputs: math.multiply_no_nan(inputs[0], inputs[1]))
     np.testing.assert_allclose(grad_fn((x, y)), (y, x), rtol=1e-3)
 
+  def test_multiply_no_nan_checkify(self):
+    """`multiply_no_nan` should not trigger checkify's NaN checks."""
+    from jax.experimental import checkify
+
+    def f(x, y):
+      return math.multiply_no_nan(x, y)
+
+    checked_f = checkify.checkify(f, errors=checkify.nan_checks)
+    err, out = checked_f(-jnp.inf, jnp.zeros(()))
+    err.throw()  # Raises if a NaN check was triggered.
+    self.assertEqual(out, 0.)
+
   def test_power_no_nan(self):
     zero = jnp.zeros(())
     nan = zero / zero