jax eig now has GPU support

Author: fchollet

keras/src/ops/linalg_test.py

@@ -350,14 +350,6 @@ def test_det(self):
     def test_eig(self):
         x = np.random.rand(2, 3, 3)
         x = x @ x.transpose((0, 2, 1))
-        if backend.backend() == "jax":
-            import jax
-
-            if jax.default_backend() == "gpu":
-                # eig not implemented for jax on gpu backend
-                with self.assertRaises(NotImplementedError):
-                    linalg.eig(x)
-                return
         w, v = map(ops.convert_to_numpy, linalg.eig(x))
         x_reconstructed = (v * w[..., None, :]) @ v.transpose((0, 2, 1))
         self.assertAllClose(x_reconstructed, x, atol=1e-4)