From fa75523767ed786e44aa5498c94fbe15834aad9d Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 30 Sep 2020 08:26:55 -0700
Subject: [PATCH] Add shape inference function for XlaScatter

PiperOrigin-RevId: 334603580
Change-Id: Idc193c32dd429cdf7f14a8496c6340c4e7b803b4
---
 tensorflow/compiler/tf2xla/ops/xla_ops.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/compiler/tf2xla/ops/xla_ops.cc b/tensorflow/compiler/tf2xla/ops/xla_ops.cc
index 2f895b17219..f73d2b109a1 100644
--- a/tensorflow/compiler/tf2xla/ops/xla_ops.cc
+++ b/tensorflow/compiler/tf2xla/ops/xla_ops.cc
@@ -762,7 +762,7 @@ REGISTER_OP("XlaScatter")
     .Attr("T: numbertype")
     .Attr("Tindices: {int32, int64}")
     .Output("output: T")
-    .SetShapeFn(UnchangedRank)
+    .SetShapeFn(shape_inference::UnchangedShape)
     .Doc(R"doc(
 Wraps the XLA Scatter operator documented at
   https://www.tensorflow.org/xla/operation_semantics#scatter.