mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 03:46:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Re-enable disabled pjit tests due to MSAN failure.

Browse Source 
PiperOrigin-RevId: 613266308
This commit is contained in:
Tongfei Guo 2024-03-06 10:33:46 -08:00 committed by jax authors
parent 026b2d207f
commit fc8dc8364e
1 changed files with 0 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tests/pjit_test.py
View File

@ -3898,8 +3898,6 @@ class ArrayPjitTest(jtu.JaxTestCase):
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)