From 25050960af883c2f131c1e90a2dd52b852ce3f09 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <aminim@google.com>
Date: Thu, 20 Jun 2019 08:58:24 -0700
Subject: [PATCH] TensorFlow operands can be broadcasted with each other and
 with the result, relax the verifier.

PiperOrigin-RevId: 254206814
---
 mlir/include/mlir/IR/OpBase.td | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index 6f199123b0f4..2598cc989d4e 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -1137,8 +1137,9 @@ class TCOpIsBroadcastableToRes<int opId, int resId> : And<[
       IsShapedTypePred>,
     SubstLeaves<"$_self", "$_op.getResult(" # resId # ")->getType()",
       IsShapedTypePred>,
-    CPred<"OpTrait::util::getBroadcastedType($_op.getOperand(" # opId # "), "
-                                            "$_op.getResult(" # resId # "))">]>;
+    CPred<"OpTrait::util::getBroadcastedType("
+              "$_op.getOperand(" # opId # ")->getType(), "
+              "$_op.getResult(" # resId # ")->getType())">]>;
 
 // Predicate to verify that all the operands at the given `indices`
 // have the same element type.