jax.nn.standardize fix NaN output.

jax/_src/nn/functions.py

@@ -693,7 +693,11 @@ def standardize(x: ArrayLike,
     # when used in neural network normalization layers
     variance = jnp.mean(
         jnp.square(x), axis, keepdims=True, where=where) - jnp.square(mean)
-  return jnp.subtract(x, jnp.asarray(mean)) * lax.rsqrt(jnp.asarray(variance) + epsilon)
+    # Because we're using a less accurate variance definition, it may return
+    # negative values. This is problematic for the rsqrt, so we take the
+    # absolute value.
+    variance = lax.abs(variance)
+  return jnp.subtract(x, mean) * lax.rsqrt(variance + epsilon)
 
 # TODO(slebedev): Change the type of `x` to `ArrayLike`.
 @api.jit(static_argnames=("num_classes", "dtype", "axis"))

tests/nn_test.py

@@ -645,6 +645,12 @@ def testStandardizeWhereMask(self):
 
     self.assertAllClose(out_masked, out_filtered)
 
+  def testStandardizeNegativeVariance(self):
+    # Regression test for https://github.com/google/jax/issues/30426
+    x = jnp.array([-11., -11., -11.]) + 3e-6
+    result = jax.nn.standardize(x)
+    self.assertFalse(jnp.any(jnp.isnan(result)))
+
   def testOneHot(self):
     actual = nn.one_hot(jnp.array([0, 1, 2]), 3)
     expected = jnp.array([[1., 0., 0.],