From 01cff3f986259d661103412a20745928c727326f Mon Sep 17 00:00:00 2001
From: Mihai Maruseac <mihaimaruseac@google.com>
Date: Mon, 2 Aug 2021 13:33:05 -0700
Subject: [PATCH] Fix heap OOB due to dimension mismatch in
 `ResourceScatterUpdate`

PiperOrigin-RevId: 388292801
Change-Id: Id9bd7244d98d41b1517d4771850b32782c0cc949
---
 tensorflow/core/kernels/resource_variable_ops.cc | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/tensorflow/core/kernels/resource_variable_ops.cc b/tensorflow/core/kernels/resource_variable_ops.cc
index b9c883c7..1c4ebb46 100644
--- a/tensorflow/core/kernels/resource_variable_ops.cc
+++ b/tensorflow/core/kernels/resource_variable_ops.cc
@@ -926,11 +926,12 @@ class ResourceScatterUpdateOp : public OpKernel {
                         params->dim_size(0), ")"));
       } else {
         int64 num_updates = updates.NumElements();
-        OP_REQUIRES(c, num_updates % N == 0,
-                    errors::InvalidArgument(
-                        "shape of indices (", indices.shape().DebugString(),
-                        ") is not compatible with the shape of updates (",
-                        updates.shape().DebugString(), ")"));
+	OP_REQUIRES(
+            c, TensorShapeUtils::StartsWith(updates.shape(), indices.shape()),
+	     errors::InvalidArgument(
+		 "The shape of indices (", indices.shape().DebugString(),
+		 ") must be a prefix of the shape of updates (",
+		 updates.shape().DebugString(), ")"));
         auto updates_flat = updates.shaped<T, 2>({N, num_updates / N});
 
         functor::ScatterFunctor<Device, T, Index, op> functor;
-- 
2.27.0