From fc8dc8364ea9c3a5704ef6f357b8db7574364bc6 Mon Sep 17 00:00:00 2001 From: Tongfei Guo Date: Wed, 6 Mar 2024 10:33:46 -0800 Subject: [PATCH] Re-enable disabled pjit tests due to MSAN failure. PiperOrigin-RevId: 613266308 --- tests/pjit_test.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/pjit_test.py b/tests/pjit_test.py index 04e9dd23dc53..5ab815f75339 100644 --- a/tests/pjit_test.py +++ b/tests/pjit_test.py @@ -3898,8 +3898,6 @@ def f(): self.assertIn('unspecified_dims=[0,1]', lowered_text) def test_jit_partially_specified_shardings(self): - if jtu.is_device_tpu(version=5, variant="e"): - self.skipTest('Remove this once b/328054509 is fixed') mesh = jtu.create_global_mesh((2, 2), ('x', 'y')) np_inp = np.arange(16).reshape(8, 2)