Adjust tolerance for lu test in preparation for XLA switch from cubla…

…s to cublasLt PiperOrigin-RevId: 474618770
google · Sep 15, 2022 · 311f85e · 311f85e
1 parent fa86ce4
commit 311f85e
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/jax/experimental/jax2tf/tests/jax2tf_limitations.py b/jax/experimental/jax2tf/tests/jax2tf_limitations.py
@@ -827,7 +827,12 @@ def _make_permutation_matrix(perm):
         custom_numeric(
             dtypes=[np.float32, np.complex64], devices=("cpu", "gpu"),
             tol=1e-5),
-        custom_numeric(dtypes=[np.float64, np.complex128], tol=1e-13),
+        custom_numeric(
+            dtypes=[np.float64, np.complex128],
+            modes=("eager", "graph"),
+            tol=1e-13),
+        custom_numeric(
+            dtypes=[np.float64, np.complex128], modes=("compiled"), tol=1e-14),
         custom_numeric(
             custom_assert=custom_assert,
             description=("May return different, but also correct, results when "