From 9cbff642515c904190511f3988cd517bb672151d Mon Sep 17 00:00:00 2001 From: Bixia Zheng Date: Tue, 28 Jan 2025 15:36:13 -0800 Subject: [PATCH] #sdy Enable test_partial_auto_of_random_keys under Shardy. PiperOrigin-RevId: 720731202 --- tests/shard_map_test.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/shard_map_test.py b/tests/shard_map_test.py index 35bb9bab1..acffbbbef 100644 --- a/tests/shard_map_test.py +++ b/tests/shard_map_test.py @@ -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)