Skip to content

Commit c098b36

Browse files
tomnatan30Google-ML-Automation
authored andcommitted
[JAX Shardy] Unskip stream annotation test when shardy is enabled, since the underlying issue is now resolved.
PiperOrigin-RevId: 738802372
1 parent 7fa7db7 commit c098b36

File tree

1 file changed

+0
-2
lines changed

1 file changed

+0
-2
lines changed

tests/memories_test.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1664,8 +1664,6 @@ class StreamAnnotationTest(jtu.JaxTestCase):
16641664
def test_stream_annotation_inside_shmap(self):
16651665
if not jtu.test_device_matches(["gpu"]):
16661666
self.skipTest("Stream annotation is only supported on GPU.")
1667-
if config.use_shardy_partitioner.value:
1668-
self.skipTest("Doesn't work with shardy")
16691667

16701668
mesh = jtu.create_mesh((2,), ('x',))
16711669
s = NamedSharding(mesh, P('x'))

0 commit comments

Comments
 (0)