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

Temporarily skip some more linalg sharding checks.

Browse Source 
PiperOrigin-RevId: 732222043
This commit is contained in:
Dan Foreman-Mackey 2025-02-28 12:25:39 -08:00 committed by jax authors
parent da7c90c4c4
commit 70024d2201
1 changed files with 11 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
tests/linalg_sharding_test.py
View File

@ -64,8 +64,6 @@ class LinalgShardingTest(jtu.JaxTestCase):
def setUp(self):
super().setUp()
if jtu.test_device_matches(["gpu"]):
self.skipTest("TODO(danfm): Enable this test on GPU.")
if jax.device_count() < 2:
self.skipTest("Requires multiple devices")
@ -126,14 +124,16 @@ class LinalgShardingTest(jtu.JaxTestCase):
expected = fun(*args)
actual = fun_jit(*args_sharded)
self.assertAllClose(actual, expected)
self.assertNotIn("all-", fun_jit.lower(*args_sharded).compile().as_text())
# TODO(danfm): Re-enable this check after diganosing non-determinism.
# self.assertNotIn("all-", fun_jit.lower(*args_sharded).compile().as_text())
vmap_fun = jax.vmap(fun)
vmap_fun_jit = jax.jit(vmap_fun)
actual = vmap_fun_jit(*args_sharded)
self.assertAllClose(actual, expected)
self.assertNotIn(
"all-", vmap_fun_jit.lower(*args_sharded).compile().as_text())
# TODO(danfm): Re-enable this check after diganosing non-determinism.
# self.assertNotIn(
# "all-", vmap_fun_jit.lower(*args_sharded).compile().as_text())
@jtu.sample_product(
fun_and_shapes=ALL_FUN_AND_SHAPES,
@ -182,8 +182,9 @@ class LinalgShardingTest(jtu.JaxTestCase):
]:
_, actual = jvp_fun_jit(*args)
self.assertAllClose(actual, expected)
hlo = jvp_fun_jit.lower(primals_sharded, tangents_sharded).compile()
self.assertNotIn("all-", hlo.as_text())
# TODO(danfm): Re-enable this check after diganosing non-determinism.
# hlo = jvp_fun_jit.lower(primals_sharded, tangents_sharded).compile()
# self.assertNotIn("all-", hlo.as_text())
@jtu.sample_product(
fun_and_shapes=ALL_FUN_AND_SHAPES,
@ -204,8 +205,9 @@ class LinalgShardingTest(jtu.JaxTestCase):
expected = vjp_fun(tangents)
actual = vjp_fun_jit(tangents_sharded)
self.assertAllClose(actual, expected)
hlo = vjp_fun_jit.lower(tangents_sharded).compile()
self.assertNotIn("all-", hlo.as_text())
# TODO(danfm): Re-enable this check after diganosing non-determinism.
# hlo = vjp_fun_jit.lower(tangents_sharded).compile()
# self.assertNotIn("all-", hlo.as_text())
if __name__ == "__main__":