Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix heap OOB due to dimension mismatch in ResourceScatterUpdate
PiperOrigin-RevId: 388292801
Change-Id: Id9bd7244d98d41b1517d4771850b32782c0cc949
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Aug 2, 2021
1 parent 96f364a commit 01cff3f
Showing 1 changed file with 6 additions and 5 deletions.
11 changes: 6 additions & 5 deletions tensorflow/core/kernels/resource_variable_ops.cc
Expand Up @@ -955,11 +955,12 @@ class ResourceScatterUpdateOp : public OpKernel {
params->dim_size(0), ")"));
} else {
int64_t 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;
Expand Down

0 comments on commit 01cff3f

Please sign in to comment.