Merge pull request #51975 from yongtang:51936-max_pool3d

tensorflower-gardener · mihaimaruseac · commit eef5b3625223 · 2021-10-28T15:22:59.000-07:00
PiperOrigin-RevId: 401245519
Change-Id: I67d2cbb0e21729b94186ca9bf82450ff93132ff2
diff --git a/tensorflow/core/kernels/pooling_ops_3d.cc b/tensorflow/core/kernels/pooling_ops_3d.cc
@@ -141,6 +141,11 @@ class Pooling3DOp : public UnaryOp<T> {
     OP_REQUIRES(context, ksize_.size() == 5,
                 errors::InvalidArgument("Sliding window ksize field must "
                                         "specify 5 dimensions"));
+    bool non_negative =
+        std::all_of(ksize_.begin(), ksize_.end(), [](int k) { return k > 0; });
+    OP_REQUIRES(context, non_negative,
+                errors::InvalidArgument("Sliding window ksize field must "
+                                        "have non-negative dimensions"));
     OP_REQUIRES_OK(context, context->GetAttr("strides", &stride_));
     OP_REQUIRES(context, stride_.size() == 5,
                 errors::InvalidArgument("Sliding window stride field must "
diff --git a/tensorflow/python/kernel_tests/pooling_ops_3d_test.py b/tensorflow/python/kernel_tests/pooling_ops_3d_test.py
@@ -510,6 +510,19 @@ def testAvgPoolGradSamePadding3_1_3d(self):
         strides=(1, 1, 1),
         padding="SAME")
 
+  def testMaxPool3DZeroPoolSize(self):
+    # Test case for GitHub issue 51936.
+    for f in [nn_ops.max_pool3d, nn_ops.avg_pool3d]:
+      with self.session():
+        with self.assertRaises((errors.InvalidArgumentError, ValueError)):
+          input_sizes = [3, 4, 10, 11, 12]
+
+          input_data = 1.
+          input_tensor = constant_op.constant(
+              input_data, shape=input_sizes, name="input")
+          pool_3d = f(input_tensor, ksize=[2, 2, 0], strides=1, padding="VALID")
+          self.evaluate(pool_3d)
+
 
   def testMaxPoolGradEagerShapeErrors(self):
     with context.eager_mode():
