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

#sdy Enable test_partial_auto_of_random_keys under Shardy.

Browse Source 
PiperOrigin-RevId: 720731202
This commit is contained in:
Bixia Zheng 2025-01-28 15:36:13 -08:00 committed by jax authors
parent 23c8607bab
commit 9cbff64251
1 changed files with 0 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tests/shard_map_test.py
View File

@ -2197,9 +2197,6 @@ class ShardMapTest(jtu.JaxTestCase):
# f(x) # don't crash
def test_partial_auto_of_random_keys(self):
if config.use_shardy_partitioner.value:
self.skipTest('Shardy does not support full-to-shard.')
mesh = jtu.create_mesh((4, 2), ('i', 'j'))
keys = jax.random.split(jax.random.key(0), 8)