[pmap] Add ignore warnings for tests failing on PmapSharding deprecation warnings.

See: https://github.com/jax-ml/jax/actions/runs/19429155535/job/55583665289
PiperOrigin-RevId: 833346834

Author: danielsuo

tests/multiprocess/host_callback_test.py

@@ -111,6 +111,7 @@ def callback_func(axis_index, x, sum_global):
                     10 * local_device_idx + 11], dtype=np.int32),
           expected_sum_global)])
 
+  @jtu.ignore_warning(category=DeprecationWarning)
   def test_io_callback_pjit(self):
     devices = np.array(sorted_devices()).reshape(
         (NR_PROCESSES, NR_LOCAL_DEVICES))

tests/multiprocess/pjit_test.py

@@ -90,6 +90,7 @@ def create_2d_non_contiguous_mesh2():
 # TODO(apaszke): Test with mesh that has host-tiled axes (especially nesting!)
 class PJitTestMultiHost(jt_multiprocess.MultiProcessTest):
 
+  @jtu.ignore_warning(category=DeprecationWarning)
   def testLocalInputsWithJaxArray(self):
     # Note that this is too small to shard over the global mesh, but fine for
     # the local mesh and so should be accepted.