fix minor issues

openxla · Nov 19, 2024 · c9202fa · c9202fa
1 parent d3657f8
commit c9202fa
Show file tree

Hide file tree

Showing 5 changed files with 23 additions and 24 deletions.
diff --git a/xla/hlo/builder/xla_builder.cc b/xla/hlo/builder/xla_builder.cc
@@ -4093,9 +4093,9 @@ XlaOp XlaBuilder::CollectivePermuteImpl(
   return ReportErrorOrReturn([&]() -> absl::StatusOr<XlaOp> {
     TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand));
     HloInstructionProto instr;
-    TF_ASSIGN_OR_RETURN(Shape shape,
-                        ShapeInference::InferCollectivePermuteShape(
-                            {{operand_shape}}, inplace));
+    TF_ASSIGN_OR_RETURN(
+        Shape shape,
+        ShapeInference::InferCollectivePermuteShape(operand_shape, inplace));
     *instr.mutable_shape() = shape.ToProto();
 
     for (const auto& pair : source_target_pairs) {
@@ -4127,11 +4127,13 @@ XlaOp XlaBuilder::CollectivePermuteImpl(
       TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand));
       operand_shapes.push_back(operand_shape);
     }
-    CHECK(operand_shapes.size() > 1);
+    CHECK_GT(operand_shapes.size(), 1);
     HloInstructionProto instr;
-    TF_ASSIGN_OR_RETURN(
-        Shape shape,
-        ShapeInference::InferCollectivePermuteShape(operand_shapes, inplace));
+    auto tuple_operand_shapes =
+        ShapeUtil::MakeTupleShapeWithPtrs(operand_shapes);
+    TF_ASSIGN_OR_RETURN(Shape shape,
+                        ShapeInference::InferCollectivePermuteShape(
+                            &tuple_operand_shapes, inplace));
     *instr.mutable_shape() =
         ShapeUtil::MakeTupleShapeWithPtrs(operand_shapes).ToProto();
 

diff --git a/xla/python/ops.cc b/xla/python/ops.cc
@@ -798,6 +798,6 @@ void BuildOpsSubmodule(nb::module_& m) {
   UNARY_OP(Conj);
   UNARY_OP(OptimizationBarrier);
 #undef UNARY_OP
-}
+}  // NOLINT(readability/fn_size)
 
 }  // namespace xla
diff --git a/xla/service/hlo_verifier.cc b/xla/service/hlo_verifier.cc
@@ -857,9 +857,10 @@ absl::Status ShapeVerifier::HandleCollectivePermute(HloInstruction* hlo) {
   absl::c_transform(
       hlo->operands(), std::back_inserter(operand_shapes),
       [](const HloInstruction* operand) { return &(operand->shape()); });
+  auto tuple_operand_shapes = ShapeUtil::MakeTupleShapeWithPtrs(operand_shapes);
   return CheckShape(hlo,
                     ShapeInference::InferCollectivePermuteShape(
-                        operand_shapes,
+                        &tuple_operand_shapes,
                         Cast<HloCollectivePermuteInstruction>(hlo)->inplace()));
 }
 

diff --git a/xla/service/shape_inference.cc b/xla/service/shape_inference.cc
@@ -2723,7 +2723,16 @@ ShapeInference::InferCollectiveBroadcastShape(
 }
 
 /* static */ absl::StatusOr<Shape> ShapeInference::InferCollectivePermuteShape(
-    absl::Span<const Shape* const> operand_shapes, bool inplace) {
+    const Shape* const operand_shape, bool inplace) {
+  std::vector<const Shape*> operand_shapes;
+  if (!operand_shape->IsTuple()) {
+    operand_shapes = {operand_shape};
+  } else {
+    absl::c_transform(operand_shape->tuple_shapes(),
+                      std::back_inserter(operand_shapes),
+                      [](const Shape& shape) { return &shape; });
+  }
+  CHECK_GT(operand_shapes.size(), 0);
   if (!inplace) {
     for (const Shape* operand_shape : operand_shapes) {
       TF_RETURN_IF_ERROR(
@@ -2739,17 +2748,6 @@ ShapeInference::InferCollectiveBroadcastShape(
   }
 }
 
-/* static */ absl::StatusOr<Shape> ShapeInference::InferCollectivePermuteShape(
-    const Shape* const operand_shape, bool inplace) {
-  CHECK(operand_shape->IsTuple());
-  std::vector<const Shape*> operand_shapes;
-  absl::c_transform(operand_shape->tuple_shapes(),
-                    std::back_inserter(operand_shapes),
-                    [](const Shape& shape) { return &shape; });
-  CHECK(operand_shapes.size() > 1);
-  return InferCollectivePermuteShape(operand_shapes, inplace);
-}
-
 /* static */ absl::StatusOr<Shape>
 ShapeInference::InferCollectivePermuteStartShape(
     absl::Span<const Shape* const> operand_shapes,
@@ -2782,7 +2780,7 @@ ShapeInference::InferCollectivePermuteStartShape(
   absl::c_transform(operand_shape->tuple_shapes(),
                     std::back_inserter(operand_shapes),
                     [](const Shape& shape) { return &shape; });
-  CHECK(operand_shapes.size() > 1);
+  CHECK_GT(operand_shapes.size(), 1);
   return InferCollectivePermuteStartShape(operand_shapes, context_shapes,
                                           inplace);
 }

diff --git a/xla/service/shape_inference.h b/xla/service/shape_inference.h
@@ -187,8 +187,6 @@ class ShapeInference {
       absl::Span<const Shape* const> operand_shapes);
 
   // Infers the shape of a collective permute operation.
-  static absl::StatusOr<Shape> InferCollectivePermuteShape(
-      absl::Span<const Shape* const> operand_shapes, bool inplace = false);
   static absl::StatusOr<Shape> InferCollectivePermuteShape(
       const Shape* const operand_shape, bool inplace = false);