[Fix] Fix the auto-sharding bug of gather (#787)

alpa-projects · Nov 27, 2022 · 181de4f · 181de4f
1 parent a194f59
commit 181de4f
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 6 deletions.
diff --git a/tests/shard_parallel/test_basic.py b/tests/shard_parallel/test_basic.py
@@ -116,7 +116,6 @@ def loss_func(params):
         n_total, n_allreduce, _, _, _ = count_communication_primitives(hlo_ir)
         assert n_total == n_allreduce == 1
 
-    @unittest.skip("The support of Gather is broken after a rebase.")
     def test_gather(self):
 
         class Model(nn.Module):
@@ -153,10 +152,10 @@ def loss_func(params):
         assert executable.auto_sharding_objective < 1e6
 
         hlo_ir = executable.get_hlo_text()
-        assert "gather(f32[64,32]" in hlo_ir
-        assert "scatter(f32[64,32]" in hlo_ir
-        _, n_allreduce, _, _, _ = count_communication_primitives(hlo_ir)
-        assert n_allreduce == 1
+        assert "gather(f32[64,32]" in hlo_ir or "gather(f32[32,64]" in hlo_ir
+        assert "scatter(f32[64,32]" in hlo_ir or "scatter(f32[32,64]" in hlo_ir
+        n_total, n_allreduce, _, _, _ = count_communication_primitives(hlo_ir)
+        assert n_total == n_allreduce == 1
 
     def test_reshape_uneven_partition(self):
         # TODO(lmzheng): Support the uneven partition of reshape.

diff --git a/third_party/tensorflow-alpa b/third_party/tensorflow-alpa