Enable some more C API tests.

PiperOrigin-RevId: 627065492

Author: Jieying Luo

tests/pytorch_interoperability_test.py

@@ -88,8 +88,6 @@ def testJaxToTorch(self, shape, dtype):
 
   @jtu.sample_product(shape=all_shapes, dtype=torch_dtypes)
   def testJaxArrayToTorch(self, shape, dtype):
-    if xla_bridge.using_pjrt_c_api():
-      self.skipTest("DLPack support is incomplete in the PJRT C API")
     if not config.enable_x64.value and dtype in [
         jnp.int64,
         jnp.float64,
@@ -111,8 +109,6 @@ def testJaxArrayToTorch(self, shape, dtype):
         self.assertAllClose(np, y.cpu().numpy())
 
   def testTorchToJaxInt64(self):
-    if xla_bridge.using_pjrt_c_api():
-      self.skipTest("DLPack support is incomplete in the PJRT C API")
     # See https://github.com/google/jax/issues/11895
     x = jax.dlpack.from_dlpack(
         torch.utils.dlpack.to_dlpack(torch.ones((2, 3), dtype=torch.int64)))
@@ -121,8 +117,6 @@ def testTorchToJaxInt64(self):
 
   @jtu.sample_product(shape=all_shapes, dtype=torch_dtypes)
   def testTorchToJax(self, shape, dtype):
-    if xla_bridge.using_pjrt_c_api():
-      self.skipTest("DLPack support is incomplete in the PJRT C API")
     if not config.enable_x64.value and dtype in [
         jnp.int64,
         jnp.float64,

tests/shard_map_test.py

@@ -1991,8 +1991,6 @@ class CustomPartitionerTest(jtu.JaxTestCase):
   def skip_if_custom_partitioning_not_supported(self):
     if jtu.is_cloud_tpu():
       raise unittest.SkipTest("Custom partitioning is not supported on libtpu.")
-    if xla_bridge.using_pjrt_c_api():
-      raise unittest.SkipTest('custom partitioning not implemented in PJRT C API')
 
   def test_custom_partitioning(self):
     self.skip_if_custom_partitioning_not_supported()