From 01cff3f986259d661103412a20745928c727326f Mon Sep 17 00:00:00 2001 From: Mihai Maruseac 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({N, num_updates / N}); functor::ScatterFunctor functor; -- 2.27.0