48 lines
1.9 KiB
Diff
48 lines
1.9 KiB
Diff
From ac117ee8a8ea57b73d34665cdf00ef3303bc0b11 Mon Sep 17 00:00:00 2001
|
|
From: Mihai Maruseac <mihaimaruseac@google.com>
|
|
Date: Fri, 30 Jul 2021 22:23:28 -0700
|
|
Subject: [PATCH] Prevent division by 0 in `resource_variable_ops.cc`
|
|
|
|
PiperOrigin-RevId: 387939939
|
|
Change-Id: Ib04902d63756633999959a70613f2eaa30c2c151
|
|
---
|
|
tensorflow/core/kernels/resource_variable_ops.cc | 11 +++++++++--
|
|
1 file changed, 9 insertions(+), 2 deletions(-)
|
|
|
|
diff --git a/tensorflow/core/kernels/resource_variable_ops.cc b/tensorflow/core/kernels/resource_variable_ops.cc
|
|
index b9c883c7..f6d114a8 100644
|
|
--- a/tensorflow/core/kernels/resource_variable_ops.cc
|
|
+++ b/tensorflow/core/kernels/resource_variable_ops.cc
|
|
@@ -688,7 +688,8 @@ class ResourceGatherOp : public OpKernel {
|
|
copy_functor(c->eigen_device<Device>(), tmp_indices.flat<Index>(),
|
|
indices.flat<Index>());
|
|
|
|
- AddBatchOffsets(&tmp_indices, params);
|
|
+ AddBatchOffsets(c, &tmp_indices, params);
|
|
+ if (!c->status().ok()) return;
|
|
op_indices = &tmp_indices;
|
|
}
|
|
|
|
@@ -720,11 +721,17 @@ class ResourceGatherOp : public OpKernel {
|
|
// Example: batch_dims = 1, indices = [[0, 1, 2], [0, 1, 2]]
|
|
// If indexing into a params dimension of size 4, then the indices will become
|
|
// [0, 1, 2, 4, 5, 6]
|
|
- void AddBatchOffsets(Tensor* indices, const Tensor& params) {
|
|
+ void AddBatchOffsets(OpKernelContext* ctx, Tensor* indices,
|
|
+ const Tensor& params) {
|
|
int64 batch_size = 1; // The size of all batch dimensions.
|
|
for (int idx = 0; idx < batch_dims_; ++idx) {
|
|
batch_size *= params.dim_size(idx);
|
|
}
|
|
+ OP_REQUIRES(
|
|
+ ctx, batch_size != 0,
|
|
+ errors::InvalidArgument(
|
|
+ "Inner size of indices would result in batch_size of 0 and a ",
|
|
+ "division by 0 in the implementation. This is illegal"));
|
|
|
|
auto indices_flat = indices->flat<Index>();
|
|
int64 const index_inner_size = indices->NumElements() / batch_size;
|
|
--
|
|
2.27.0
|
|
|