Relax test tolerance for complex128 pow in lax_test.py.

This is failing in CI in some CPU configurations. PiperOrigin-RevId: 705558897
jax-ml · Dec 12, 2024 · 6548caf · 6548caf
1 parent ea63aea
commit 6548caf
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/lax_test.py b/tests/lax_test.py
@@ -129,6 +129,8 @@ def testOpAgainstNumpy(self, op_name, rng_factory, shapes, dtype, tol):
         tol = jtu.join_tolerance(tol, 1e-3)
       elif op_name == "lgamma" and dtype == np.float32:
         tol = jtu.join_tolerance(tol, 1e-3)
+    elif op_name == "pow" and dtype == np.complex128:
+      tol = jtu.join_tolerance(tol, 2e-15)
     self._CheckAgainstNumpy(numpy_op, op, args_maker, tol=tol)
 
   # TODO test shift_left, shift_right_arithmetic, shift_right_logical