From 5fa4d076f48d6e073bf0cd457eb1db94f003d691 Mon Sep 17 00:00:00 2001
From: Tongfei Guo <tongfei@google.com>
Date: Mon, 8 Apr 2024 22:36:47 -0700
Subject: [PATCH] Allow setting operand_shapes_with_layout on custom call
 instruction.

PiperOrigin-RevId: 623055564
---
 third_party/xla/xla/hlo/ir/hlo_instructions.h | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.h b/third_party/xla/xla/hlo/ir/hlo_instructions.h
index e706f80c971f33..f11e536f45e675 100644
--- 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;
   }