From 136b51f10903e044308cf77117c0ed9871350475 Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Fri, 30 Jul 2021 20:50:00 -0700 Subject: [PATCH] Add missing validation to `maxpooling_op.cc` PiperOrigin-RevId: 387932441 Change-Id: I43a0b24e6a12cc965611144ba035accd384594b9 --- tensorflow/core/kernels/maxpooling_op.cc | 5 +++++ tensorflow/core/kernels/pooling_ops_common.cc | 2 ++ 2 files changed, 7 insertions(+) diff --git a/tensorflow/core/kernels/maxpooling_op.cc b/tensorflow/core/kernels/maxpooling_op.cc index 7accd1a8..27ff9005 100644 --- a/tensorflow/core/kernels/maxpooling_op.cc +++ b/tensorflow/core/kernels/maxpooling_op.cc @@ -68,6 +68,7 @@ static void SpatialMaxPoolWithArgMaxHelper( "SpatialMaxPoolWithArgMaxHelper requires include_batch_in_index " "to be True when when input_backprop != nullptr")); } + if (tensor_in.NumElements() == 0 || output->NumElements() == 0) return; typedef Eigen::Map> ConstEigenMatrixMap; @@ -783,6 +784,10 @@ class MaxPoolingNoMaskOp : public OpKernel { void Compute(OpKernelContext* context) override { const Tensor& tensor_in = context->input(0); + OP_REQUIRES(context, tensor_in.dims() == 4, + errors::InvalidArgument("tensor_in must be 4-dimensional (2)")); + OP_REQUIRES(context, tensor_in.NumElements() > 0, + errors::InvalidArgument("tensor_in must not be empty (2)")); PoolParameters params{context, ksize_, stride_, padding_, data_format_, tensor_in.shape()}; diff --git a/tensorflow/core/kernels/pooling_ops_common.cc b/tensorflow/core/kernels/pooling_ops_common.cc index 4bd71054..2af93960 100644 --- a/tensorflow/core/kernels/pooling_ops_common.cc +++ b/tensorflow/core/kernels/pooling_ops_common.cc @@ -96,6 +96,8 @@ PoolParameters::PoolParameters(OpKernelContext* context, pad_depth = 0; out_depth = depth; } else { + OP_REQUIRES(context, depth_window > 0, + errors::InvalidArgument("depth_window must not be 0")); // Our current version of depthwise max pooling does not support // any padding, and expects the depth_window to equal the // depth_stride (no overlapping). -- 2.27.0