Skip to content
Permalink
Browse files Browse the repository at this point in the history
Add missing validation to maxpooling_op.cc
PiperOrigin-RevId: 387932441
Change-Id: I43a0b24e6a12cc965611144ba035accd384594b9
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Jul 31, 2021
1 parent 1ef3a71 commit 136b51f
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 0 deletions.
5 changes: 5 additions & 0 deletions tensorflow/core/kernels/maxpooling_op.cc
Expand Up @@ -74,6 +74,7 @@ static void SpatialMaxPoolWithArgMaxHelper(
errors::Internal("SpatialMaxPoolWithArgMaxHelper requires Targmax "
"to be int64 when input_backprop != nullptr"));
}
if (tensor_in.NumElements() == 0 || output->NumElements() == 0) return;

typedef Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>
ConstEigenMatrixMap;
Expand Down Expand Up @@ -949,6 +950,10 @@ class MaxPoolingWithArgmaxOp : 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_,
Expand Down
2 changes: 2 additions & 0 deletions tensorflow/core/kernels/pooling_ops_common.cc
Expand Up @@ -171,6 +171,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).
Expand Down

0 comments on commit 136b51f

Please sign in to comment.