Fix back compat test to ignore warnings

PiperOrigin-RevId: 842272730

Author: yashk2810

jax/_src/test_util.py

@@ -1248,8 +1248,6 @@ class JaxTestCase(parameterized.TestCase):
     'jax_legacy_prng_key': 'error',
   }
 
-
-
   def setUp(self):
     super().setUp()
     self.enterContext(assert_global_configs_unchanged())

tests/export_back_compat_test.py

@@ -784,6 +784,8 @@ 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:
@@ -1006,8 +1008,10 @@ def check_top_k_results(res_run, res_expected, *, rtol, atol):
     )
 
 
-@jtu.with_config(jax_use_shardy_partitioner=True)
 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.
@@ -1046,7 +1050,5 @@ def shard_map_func(x):  # b: f32[2, 4]
             expect_current_custom_calls=custom_call_targets_override)
 
 
-
-
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())

tests/fused_attention_stablehlo_test.py

@@ -264,6 +264,8 @@ def dot_product_attention_fp8(query, key, value, fp8_metas):
   return out[0], (query_grad, key_grad, value_grad)
 
 
+@jtu.ignore_warning(category=DeprecationWarning,
+                    message='`with mesh:` context manager')
 class DotProductAttentionTest(jtu.JaxTestCase):
   def setUp(self):
     super().setUp()