Re-enable disabled pjit tests due to MSAN failure.

PiperOrigin-RevId: 613266308
google · Mar 6, 2024 · fc8dc83 · fc8dc83
1 parent 026b2d2
commit fc8dc83
Showing 1 changed file with 0 additions and 2 deletions.
diff --git 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)