From 152cf2c46e8559504d5dafd74d1e085517f93ba4 Mon Sep 17 00:00:00 2001
From: Isha Arkatkar <ishark@google.com>
Date: Wed, 24 Nov 2021 16:03:38 -0800
Subject: [PATCH] Fix potential divide by zero error when executing
 FractionalMaxPool, when pooling ratio is higher than input size for a
 particular dimension.

PiperOrigin-RevId: 412151722
Change-Id: I06e57cbb8eca43816eff79eac264fa7aae8f7163
---
 .../core/kernels/fractional_max_pool_op.cc      |  7 +++++++
 .../kernel_tests/fractional_max_pool_op_test.py | 17 +++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/tensorflow/core/kernels/fractional_max_pool_op.cc b/tensorflow/core/kernels/fractional_max_pool_op.cc
index 7519c84667409f..0722c408fba9d4 100644
--- a/tensorflow/core/kernels/fractional_max_pool_op.cc
+++ b/tensorflow/core/kernels/fractional_max_pool_op.cc
@@ -83,6 +83,13 @@ class FractionalMaxPoolOp : public OpKernel {
     std::vector<int> output_size(tensor_in_and_out_dims);
     for (int i = 0; i < tensor_in_and_out_dims; ++i) {
       input_size[i] = tensor_in.dim_size(i);
+
+      OP_REQUIRES(
+          context, input_size[i] >= pooling_ratio_[i],
+          errors::InvalidArgument("Pooling ratio is higher than input "
+                                  "dimension size for dimension ",
+                                  i, ". Input dim size: ", input_size[i],
+                                  " pooling ratio: ", pooling_ratio_[i]));
     }
     // Output size.
     for (int i = 0; i < tensor_in_and_out_dims; ++i) {
diff --git a/tensorflow/python/kernel_tests/fractional_max_pool_op_test.py b/tensorflow/python/kernel_tests/fractional_max_pool_op_test.py
index 2b1e30a8bbd606..f395fefeb22f87 100644
--- a/tensorflow/python/kernel_tests/fractional_max_pool_op_test.py
+++ b/tensorflow/python/kernel_tests/fractional_max_pool_op_test.py
@@ -24,6 +24,7 @@
 
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import gen_nn_ops
@@ -307,6 +308,22 @@ def testDifferentInputTensorShape(self):
           input_b, row_seq, col_seq, overlapping)
       self.assertSequenceEqual(expected.shape, actual.shape)
 
+  def testDeterminismExceptionThrowing(self):
+    tensor_shape = (5, 20, 20, 3)
+    rand_mat = self._PRNG.random_sample(tensor_shape) * 1000 - 500
+    with test_util.deterministic_ops():
+      with self.assertRaisesRegex(
+          ValueError, "requires a non-zero seed to be passed in when "
+          "determinism is enabled"):
+        nn_ops.fractional_max_pool_v2(rand_mat, [1, 1.5, 1.5, 1])
+      nn_ops.fractional_max_pool_v2(rand_mat, [1, 1.5, 1.5, 1], seed=1)
+
+      with self.assertRaisesRegex(ValueError,
+                                  'requires "seed" and "seed2" to be non-zero'):
+        nn_ops.fractional_max_pool(rand_mat, [1, 1.5, 1.5, 1])
+      nn_ops.fractional_max_pool(
+          rand_mat, [1, 1.5, 1.5, 1], seed=1, seed2=1, deterministic=True)
+
 
 class FractionalMaxPoolGradTest(test.TestCase):
   """Tests for FractionalMaxPoolGrad.