Some cleanups

pytorch · Nov 29, 2023 · 8e6715e · 8e6715e
1 parent 5843b43
commit 8e6715e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 7 deletions.
diff --git a/torch_xla/csrc/cross_replica_reduces.cpp b/torch_xla/csrc/cross_replica_reduces.cpp
@@ -235,7 +235,7 @@ AllGatherResult BuildAllGather(absl::Span<const xla::XlaOp> inputs,
           xla::AllGather(xla::Tuple(inputs[0].builder(), type_ctx.second.ops),
                          dim, shard_count, cc_groups);
     }
-    if (type_ctx.second.indices.size() > 1) {
+    if (ShapeHelper::ShapeOfXlaOp(all_gather_result).rank() == 0) {
       for (size_t i = 0; i < type_ctx.second.indices.size(); ++i) {
         size_t op_idx = type_ctx.second.indices[i];
         result[op_idx] = xla::GetTupleElement(all_gather_result, i);

diff --git a/torch_xla/csrc/ops/all_gather.cpp b/torch_xla/csrc/ops/all_gather.cpp
@@ -19,12 +19,8 @@ xla::Shape NodeOutputShape(c10::ArrayRef<torch::lazy::Value> inputs,
     AllGatherResult result =
         BuildAllGather(operands.subspan(0, operands.size() - 1),
                        operands.back(), dim, shard_count, groups, pin_layout);
-    std::vector<xla::XlaOp> outputs;
-    for (size_t i = 0; i < result.result.size(); ++i) {
-      outputs.emplace_back(result.result[i]);
-    }
-    outputs.emplace_back(result.token);
-    return xla::Tuple(operands[0].builder(), outputs);
+    result.result.emplace_back(result.token);
+    return xla::Tuple(operands[0].builder(), result.result);
   };
   std::vector<xla::Shape> input_shapes;
   for (const auto& input : inputs) {