38 lines
1.6 KiB
Diff
38 lines
1.6 KiB
Diff
|
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
|
||
|
|
|