From a58e62590e0d1bde5a139bc0f6b243578b038ee3 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar@google.com>
Date: Tue, 16 Mar 2021 16:53:10 -0700
Subject: [PATCH] Restrict canonicalization to avoid changing type

Issue #47516

PiperOrigin-RevId: 363300979
---
 lib/Dialect/mhlo/IR/hlo_patterns.td | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/lib/Dialect/mhlo/IR/hlo_patterns.td b/lib/Dialect/mhlo/IR/hlo_patterns.td
index 73fca2d..58193a8 100644
--- a/lib/Dialect/mhlo/IR/hlo_patterns.td
+++ b/lib/Dialect/mhlo/IR/hlo_patterns.td
@@ -36,12 +36,13 @@ def DynamicBroadcastToOwnShape_4 : Pat<
   (HLO_DynamicBroadcastInDimOp:$op $x, (Tensor_CastOp (Shape_ShapeOfOp $x)), $attr),
   (Tensor_CastOp $x)>;
 
-def ShapeOfDynamicReshape : Pat<
-  (Shape_ShapeOfOp (HLO_DynamicReshapeOp $x, $shape)),
-  (replaceWithValue $shape)>;
-
 def HasSameType : Constraint<CPred<"$0.getType() == $1.getType()">>;
 
+def ShapeOfDynamicReshape : Pat<
+  (Shape_ShapeOfOp:$op (HLO_DynamicReshapeOp $x, $shape)),
+  (replaceWithValue $shape),
+  [(HasSameType $shape, $op)]>;
+
 def IdentityBroadcastReshape : Pat<
   (HLO_ReshapeOp:$op (HLO_BroadcastOp $input, $dims)),
   (replaceWithValue $input),