Ignore some linesearch tests on gpu/tpu

PiperOrigin-RevId: 643364817

optax/_src/linesearch_test.py

@@ -489,7 +489,7 @@ def test_failure_descent_direction(self):
     # @absltest.skipIf(jax.devices()[0] == 'tpu', reason = ...))
     # because jax arrays cannot be manipulated from the top level of the python
     # program
-    if jax.default_backend() == 'tpu':
+    if jax.default_backend() in ['tpu', 'gpu']:
       return
     else:
       # For this f and p, starting at a point on axis 0, the strong Wolfe
@@ -523,7 +523,7 @@ def test_failure_too_small_max_stepsize(self):
     # @absltest.skipIf(jax.devices()[0] == 'tpu', reason = ...))
     # because jax arrays cannot be manipulated from the top level of the python
     # program
-    if jax.default_backend() == 'tpu':
+    if jax.default_backend() in ['tpu', 'gpu']:
       return
     else:
       def fn(x):
@@ -558,7 +558,7 @@ def test_failure_not_enough_iter(self):
     # @absltest.skipIf(jax.devices()[0] == 'tpu', reason = ...))
     # because jax arrays cannot be manipulated from the top level of the python
     # program
-    if jax.default_backend() == 'tpu':
+    if jax.default_backend() in ['tpu', 'gpu']:
       return
     else:
       def fn(x):
@@ -608,7 +608,7 @@ def test_failure_flat_fun(self):
     # @absltest.skipIf(jax.devices()[0] == 'tpu', reason = ...))
     # because jax arrays cannot be manipulated from the top level of the python
     # program
-    if jax.default_backend() == 'tpu':
+    if jax.default_backend() in ['tpu', 'gpu']:
       return
     else:
       def fun_flat(x):
@@ -631,7 +631,7 @@ def test_failure_inf_value(self):
     # @absltest.skipIf(jax.devices()[0] == 'tpu', reason = ...))
     # because jax arrays cannot be manipulated from the top level of the python
     # program
-    if jax.default_backend() == 'tpu':
+    if jax.default_backend() in ['tpu', 'gpu']:
       return
     else:
       def fun_inf(x):