From e31695e95fad5995ec95a5b179e69b6a45d8c63e Mon Sep 17 00:00:00 2001
From: Blake Hechtman <blakehechtman@google.com>
Date: Mon, 28 Sep 2020 11:54:01 -0700
Subject: [PATCH] [XLA] No-op reshapes should not return true for inserting or
 deleting one-sized dimensions.

PiperOrigin-RevId: 334206368
Change-Id: I4ba86171812aafcaf9ece282ac05bdd01e914fb1
---
 tensorflow/compiler/xla/shape_util.cc      | 3 ++-
 tensorflow/compiler/xla/shape_util_test.cc | 2 ++
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/tensorflow/compiler/xla/shape_util.cc b/tensorflow/compiler/xla/shape_util.cc
index 0833919b124..c9f37bdc430 100644
--- a/tensorflow/compiler/xla/shape_util.cc
+++ b/tensorflow/compiler/xla/shape_util.cc
@@ -1074,7 +1074,8 @@ ShapeUtil::InsertedOrDeleted1SizedDimensions(const Shape& shape_pre,
     }
   }
 
-  return std::make_tuple(true, deleted_indices, inserted_indices);
+  return std::make_tuple(!deleted_indices.empty() || !inserted_indices.empty(),
+                         deleted_indices, inserted_indices);
 }
 
 /* static */ std::vector<std::pair<int64, int64>>
diff --git a/tensorflow/compiler/xla/shape_util_test.cc b/tensorflow/compiler/xla/shape_util_test.cc
index 4e2030667ee..414b53d4f67 100644
--- a/tensorflow/compiler/xla/shape_util_test.cc
+++ b/tensorflow/compiler/xla/shape_util_test.cc
@@ -558,6 +558,8 @@ TEST(ShapeUtilTest, InsertedOrDeleted1SizedDimensions) {
       ShapeUtil::InsertedOrDeleted1SizedDimensions(shape0, shape1)));
   EXPECT_FALSE(std::get<0>(
       ShapeUtil::InsertedOrDeleted1SizedDimensions(shape0, shape2)));
+  EXPECT_FALSE(std::get<0>(
+      ShapeUtil::InsertedOrDeleted1SizedDimensions(shape0, shape0)));
 }
 
 TEST(ShapeUtilTest, ForEachIndex) {