Relax some test tolerances in for_loop_test.py.

This PR attempts to fix some CI failures on Mac ARM. PiperOrigin-RevId: 672312564
jax-ml · Sep 8, 2024 · b6abd73 · b6abd73
1 parent 5af1efb
commit b6abd73
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/for_loop_test.py b/tests/for_loop_test.py
@@ -345,7 +345,8 @@ def g(a, b):
     expected_tangents = g_lin(a, b)
     _, actual_tangents = jax.jvp(g, (a, b), (a, b))
     np.testing.assert_allclose(actual_tangents[0], expected_tangents[0])
-    np.testing.assert_allclose(actual_tangents[1], expected_tangents[1])
+    np.testing.assert_allclose(actual_tangents[1], expected_tangents[1],
+                               rtol=1e-6)
 
   @jtu.sample_product(
     [dict(for_body_name=for_body_name, f=for_body, ref=ref,