Allow setting operand_shapes_with_layout on custom call instruction.

PiperOrigin-RevId: 623055564
tensorflow · May 8, 2024 · 5fa4d07 · 5fa4d07
1 parent cfb58b9
commit 5fa4d07
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.h b/third_party/xla/xla/hlo/ir/hlo_instructions.h
@@ -2094,6 +2094,11 @@ class HloCustomCallInstruction : public HloCallableInstruction {
     CHECK(layout_constrained());
     return operand_shapes_with_layout_;
   }
+  void set_operand_shapes_with_layout(
+      std::vector<Shape> operand_shapes_with_layout) {
+    CHECK(layout_constrained());
+    operand_shapes_with_layout_ = std::move(operand_shapes_with_layout);
+  }
   void set_custom_call_schedule(CustomCallSchedule custom_call_schedule) {
     custom_call_schedule_ = custom_call_schedule;
   }