Enable a number of tests on GPU.

In particular, pjit/xmap work on CPU these days.

PiperOrigin-RevId: 446085110

Author: hawkinsp

tests/api_test.py

@@ -2212,7 +2212,7 @@ def test_xla_computation_psum_constant(self):
     f = lambda: jax.lax.psum(1, "i")
     api.xla_computation(f, axis_env=[("i", 2)])()  # doesn't crash
 
-  @jtu.skip_on_devices("cpu", "gpu")
+  @jtu.skip_on_devices("cpu")
   @jtu.ignore_warning(message="Some donated buffers were not usable")
   def test_xla_computation_donate_argnums(self):
     api.xla_computation(lambda x: None, donate_argnums=(0,))(3)  # doesn't crash

tests/debug_nans_test.py

@@ -144,7 +144,6 @@ def testXmap(self):
           ans.block_until_ready()
 
   @jtu.ignore_warning(message=".*is an experimental.*")
-  @jtu.skip_on_devices("cpu", "gpu")
   def testPjit(self):
     if jax.device_count() < 2:
       raise SkipTest("test requires >=2 devices")

tests/jaxpr_effects_test.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import unittest
 
 from absl.testing import absltest
 from absl.testing import parameterized
@@ -190,8 +189,6 @@ def f(x):
       jax.make_jaxpr(f)(jnp.arange(jax.local_device_count()))
 
   def test_pjit_inherits_effects(self):
-    if jax.default_backend() not in {'gpu', 'tpu'}:
-      raise unittest.SkipTest("pjit only supports GPU and TPU backends")
     def f(x):
       effect_p.bind(effect='foo')
       effect_p.bind(effect='bar')

tests/xmap_test.py

@@ -212,8 +212,6 @@ class XMapTestCase(jtu.BufferDonationTestCase):
 # A mixin that enables SPMD lowering tests
 class SPMDTestMixin:
   def setUp(self):
-    if jtu.device_under_test() not in ['tpu', 'gpu']:
-      raise SkipTest
     super().setUp()
     jtu.set_spmd_lowering_flag(True)
 
@@ -223,8 +221,6 @@ def tearDown(self):
 
 class ManualSPMDTestMixin:
   def setUp(self):
-    if jtu.device_under_test() not in ['tpu', 'gpu']:
-      raise SkipTest
     if not hasattr(xla_client.OpSharding.Type, "MANUAL"):
       raise SkipTest
     super().setUp()