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

[JAX Shardy] Unskip stream annotation test when shardy is enabled, since the underlying issue is now resolved.

Browse Source 
PiperOrigin-RevId: 738802372
This commit is contained in:
Tom Natan 2025-03-20 08:01:08 -07:00 committed by jax authors
parent 7fa7db7a9f
commit c098b363fb
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/memories_test.py
View File

@ -1664,8 +1664,6 @@ class StreamAnnotationTest(jtu.JaxTestCase):
def test_stream_annotation_inside_shmap(self):
if not jtu.test_device_matches(["gpu"]):
self.skipTest("Stream annotation is only supported on GPU.")
if config.use_shardy_partitioner.value:
self.skipTest("Doesn't work with shardy")
mesh = jtu.create_mesh((2,), ('x',))
s = NamedSharding(mesh, P('x'))