Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix security vulnerability with AvgPoolGrad
PiperOrigin-RevId: 461724457
  • Loading branch information
tensorflower-gardener committed Jul 18, 2022
1 parent 8741e57 commit 3a6ac52
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 3 deletions.
6 changes: 3 additions & 3 deletions tensorflow/core/kernels/avgpooling_op.cc
Expand Up @@ -298,7 +298,7 @@ class AvgPoolingGradOp : public OpKernel {
TensorShape output_shape;
auto shape_vec = tensor_in_shape.vec<int32>();
for (int64_t i = 0; i < tensor_in_shape.NumElements(); ++i) {
output_shape.AddDim(shape_vec(i));
OP_REQUIRES_OK(context, output_shape.AddDimWithStatus(shape_vec(i)));
}
const int64_t in_rows = output_shape.dim_size(1);
const int64_t in_cols = output_shape.dim_size(2);
Expand Down Expand Up @@ -457,7 +457,7 @@ class AvgPoolingGradOp<GPUDevice, T> : public OpKernel {
TensorShape output_shape;
auto shape_vec = tensor_in_shape.vec<int32>();
for (int64_t i = 0; i < tensor_in_shape.NumElements(); ++i) {
output_shape.AddDim(shape_vec(i));
OP_REQUIRES_OK(context, output_shape.AddDimWithStatus(shape_vec(i)));
}

if (output_shape.num_elements() == 0) {
Expand Down Expand Up @@ -543,7 +543,7 @@ class AvgPoolingGradOpCustomGPUKernel : public OpKernel {
TensorShape output_shape;
auto shape_vec = tensor_in_shape.vec<int32>();
for (int64_t i = 0; i < tensor_in_shape.NumElements(); ++i) {
output_shape.AddDim(shape_vec(i));
OP_REQUIRES_OK(context, output_shape.AddDimWithStatus(shape_vec(i)));
}
if (output_shape.num_elements() == 0) {
Tensor* output = nullptr;
Expand Down
16 changes: 16 additions & 0 deletions tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
Expand Up @@ -2470,6 +2470,22 @@ def testMaxPoolGradWithArgmaxEagerShapeErrors(self):
inp, grad, argmax, ksize=[1, 1, 1, 1], strides=[1, 1, 1, 1],
padding="VALID")

def testAvgPoolGradInvalidInputShapeRaiseError(self):
with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)):
with self.cached_session():
orig_input_shape = constant_op.constant(
-536870912, shape=[4], dtype=dtypes.int32)
grad = constant_op.constant(
.0890338004362538, shape=[1, 5, 7, 1], dtype=dtypes.float64)
t = gen_nn_ops.AvgPoolGrad(
orig_input_shape=orig_input_shape,
grad=grad,
ksize=[1, 2, 2, 1],
strides=[1, 2, 2, 1],
padding="VALID",
data_format="NHWC")
self.evaluate(t)


def GetMaxPoolFwdTest(input_size, filter_size, strides, padding):

Expand Down

0 comments on commit 3a6ac52

Please sign in to comment.