[export] Fix the "with mesh" deprecation warning

Author: gnecula

tests/export_back_compat_test.py

@@ -784,8 +784,6 @@ def func(x):
       data = self.load_testdata(cuda_threefry2x32.data_2024_07_30)
       self.run_one_test(func, data)
 
-  @jtu.ignore_warning(category=DeprecationWarning,
-                      message='`with mesh:` context manager')
   def test_tpu_sharding(self):
     # Tests "Sharding", "SPMDShardToFullShape", "SPMDFullToShardShape" on TPU
     if not jtu.test_device_matches(["tpu"]) or len(jax.devices()) < 2:
@@ -815,7 +813,7 @@ def func(x):  # b: f32[2, 4]
     # the expected custom call targets for old test data that was serialized
     # with custom calls.
     for data, custom_call_targets_override in data:
-      with mesh:
+      with jax.set_mesh(mesh):
         if jax.config.jax_use_shardy_partitioner:
           self.run_one_test(
               func, self.load_testdata(data["shardy"]),
@@ -1010,8 +1008,6 @@ def check_top_k_results(res_run, res_expected, *, rtol, atol):
 
 class ShardyCompatTest(bctu.CompatTestBase):
 
-  @jtu.ignore_warning(category=DeprecationWarning,
-                      message='`with mesh:` context manager')
   def test_shardy_sharding_ops_with_different_meshes(self):
     # Tests whether we can save and load a module with meshes that have the
     # same axis sizes (and same order) but different axis names.
@@ -1044,7 +1040,7 @@ def shard_map_func(x):  # b: f32[2, 4]
     # the expected custom call targets for old test data that was serialized
     # with custom calls.
     for data, custom_call_targets_override in data:
-      with Mesh(devices, axis_names=('x')):
+      with jax.set_mesh(Mesh(devices, axis_names=('x'))):
         self.run_one_test(
             func, self.load_testdata(data),
             expect_current_custom_calls=custom_call_targets_override)