Update grad(shard_map(...)) test.

jax-ml · Dec 17, 2024 · 9395fc7 · 9395fc7
1 parent 3f4e54b
commit 9395fc7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/debug_nans_test.py b/tests/debug_nans_test.py
@@ -156,7 +156,7 @@ def f(x):
 
     with self.assertRaisesRegex(
         FloatingPointError,
-        r"Invalid value \(nan\) encountered in sharded computation.\nWhen differentiating"):
+        r"invalid value \(nan\) encountered in mul\nWhen differentiating"):
       ans, = f_vjp(jnp.ones([1]))
       ans.block_until_ready()