Disable pjit ArrayPjitTest.test_device_put_grad test on TPU v5e

PiperOrigin-RevId: 704378732

Author: kanglant

tests/pjit_test.py

@@ -3411,6 +3411,9 @@ def test_device_assignment_mismatch_apply_primitive(self):
   def test_device_put_grad(self):
     if jax.device_count() < 8:
       self.skipTest("Requires >=8 devices.")
+    if jtu.is_device_tpu(5, 'e'):
+      self.skipTest('TPU v5e does not support computations that run on a '
+                    'non-singleton subset of cores.')
 
     def _test(fun, inp, np_inp, in_s):
       out = fun(inp)