Adjust tests

keras-team · Apr 7, 2024 · 121055f · 121055f
1 parent 142ff58
commit 121055f
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/keras/ops/linalg_test.py b/keras/ops/linalg_test.py
@@ -356,6 +356,9 @@ def test_inv(self):
         )
 
     def test_lu_factor(self):
+        if self._jax_uses_gpu():
+            self.skipTest("Skipping test with JAX + GPU due to temporary error")
+
         def _pivot_matrix(pivots, n):
             p_matrix = np.eye(n)
             for i, p in enumerate(pivots):
@@ -477,6 +480,8 @@ def test_solve(self):
         self.assertAllClose(output, expected_result)
 
     def test_solve_triangular(self):
+        if self._jax_uses_gpu():
+            self.skipTest("Skipping test with JAX + GPU due to temporary error")
 
         # 2d-case
         x1 = np.array([[1, 2], [0, 5]], dtype="float32")

diff --git a/keras/ops/numpy_test.py b/keras/ops/numpy_test.py
@@ -4562,9 +4562,6 @@ def test_elementwise_unary_indexed_slices_correctness(
     def test_other_unary_symbolic_sparse_correctness(
         self, op_function, op_class, np_op, op_kwargs, input_shape
     ):
-        if self._tensorflow_uses_gpu():
-            self.skipTest("Skipping test with TF + GPU due to XLA error")
-
         x = np.random.random(input_shape)
         if op_function is knp.mean:
             x = create_indexed_slices(x)