From ac117ee8a8ea57b73d34665cdf00ef3303bc0b11 Mon Sep 17 00:00:00 2001 From: Mihai Maruseac 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(), tmp_indices.flat(), indices.flat()); - 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(); int64 const index_inner_size = indices->NumElements() / batch_size; -- 2.27.0