From e18c52cd16ab9f46d09450a0b7eff94c365fc5da Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar@google.com>
Date: Wed, 20 May 2020 07:11:58 -0700
Subject: [PATCH] Update StopGradients's overly strict verification

Subtypes may differ and we may have dynamic shapes, so don't require static
type equality.

PiperOrigin-RevId: 312475863
Change-Id: I4dd6fef11ece9e1e62560e411ca758e848964330
---
 tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
index fd24b7284c1..957ba4909a9 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
@@ -8317,8 +8317,9 @@ def TF_StackV2Op : TF_Op<"StackV2", []> {
   );
 }
 
-def TF_StopGradientOp : TF_Op<"StopGradient", [NoSideEffect, SameOperandsAndResultType]> {
-  let summary = "Stops gradient computation.";
+def TF_StopGradientOp : TF_Op<"StopGradient",
+    [NoSideEffect, TF_AllTypesMatch<["input", "output"]>]> {
+  let summary = "Stops gradient computation";
 
   let description = [{
 When executed in a graph, this op outputs its input tensor as-is.