From ffef5bdff2e0d6b078bd37f93a44f9e4be75c277 Mon Sep 17 00:00:00 2001
From: Piotr Szmelczynski <piotr.szmelczynski@intel.com>
Date: Wed, 15 Sep 2021 06:50:30 +0200
Subject: [PATCH] Avg pool bug fix (#7493)

* modify avg_pool ref impl to work with n_elements = 0

* add avg_pool backend test

* fix bug with division by 0
---
 .../ngraph/runtime/reference/avg_pool.hpp     | 17 ++++++------
 ngraph/test/backend/avg_pool.in.cpp           | 26 +++++++++++++++++++
 2 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/ngraph/core/reference/include/ngraph/runtime/reference/avg_pool.hpp b/ngraph/core/reference/include/ngraph/runtime/reference/avg_pool.hpp
index 94dfbdf1d46..f68865830af 100644
--- a/ngraph/core/reference/include/ngraph/runtime/reference/avg_pool.hpp
+++ b/ngraph/core/reference/include/ngraph/runtime/reference/avg_pool.hpp
@@ -215,16 +215,17 @@ void avg_pool(const T* arg,
             }
         }
 
-        if (n_elements == 0) {
-            throw std::runtime_error("AvgPool elements == 0, must be non-zero");
+        if (n_elements != 0) {
+            if (std::is_same<T, int8_t>::value || std::is_same<T, uint8_t>::value) {
+                out[output_transform.index(out_coord)] =
+                    static_cast<T>(std::nearbyint(static_cast<float>(result) / n_elements));
+            } else {
+                out[output_transform.index(out_coord)] = result / n_elements;
+            }
+        } else {
+            out[output_transform.index(out_coord)] = T{0};
         }
 
-        if (std::is_same<T, int8_t>::value || std::is_same<T, uint8_t>::value) {
-            out[output_transform.index(out_coord)] =
-                static_cast<T>(std::nearbyint(static_cast<float>(result) / n_elements));
-        } else {
-            out[output_transform.index(out_coord)] = result / n_elements;
-        }
         std::fesetround(old_mode);
     }
     NGRAPH_SUPPRESS_DEPRECATED_END
diff --git a/ngraph/test/backend/avg_pool.in.cpp b/ngraph/test/backend/avg_pool.in.cpp
index cd376fad62f..cf983ecf659 100644
--- a/ngraph/test/backend/avg_pool.in.cpp
+++ b/ngraph/test/backend/avg_pool.in.cpp
@@ -223,3 +223,29 @@ NGRAPH_TEST(${BACKEND_NAME}, avg_pool_2d_same_lower) {
     test_case.add_expected_output<float>(out_shape, result);
     test_case.run();
 }
+
+NGRAPH_TEST(${BACKEND_NAME}, avg_pool_2d_padding) {
+    Shape in_shape{1, 1, 3, 3};
+    Shape out_shape{1, 1, 3, 3};
+    const Strides& strides{2, 2};
+    const Shape& pads_begin{1, 1};
+    const Shape& pads_end{1, 1};
+    const Shape& kernel{2, 2};
+    const bool exclude_pad = true;
+    const op::RoundingType rounding_type = op::RoundingType::CEIL;
+    const op::PadType pad_type = op::PadType::NOTSET;
+
+    auto A = make_shared<op::Parameter>(element::f32, in_shape);
+    auto avgPool =
+        make_shared<op::v1::AvgPool>(A, strides, pads_begin, pads_end, kernel, exclude_pad, rounding_type, pad_type);
+    auto f = make_shared<Function>(avgPool, ParameterVector{A});
+
+    std::vector<float> a(1 * 1 * 3 * 3);
+    std::iota(std::begin(a), std::end(a), 1);
+    std::vector<float> result{1.0f, 2.5f, 0, 5.5f, 7.0f, 0, 0, 0, 0};
+
+    auto test_case = test::TestCase<TestEngine>(f);
+    test_case.add_input<float>({a});
+    test_case.add_expected_output<float>(out_shape, result);
+    test_case.run();
+}