Fix sum reference to handle corner cases with +-inf (#3412)

* Fix sum reference to handle corner cases with +-inf

* Review comments, and try to make Windows happy

* Update GPU unit_test.manifest

* More template grindery, to make macOS happy

src/ngraph/runtime/gpu/unit_test.manifest

@@ -100,6 +100,9 @@ all_2x2x3_eliminate_dims_1_2
 all_2x2x3_eliminate_dims_0_1_2
 all_dynamic
 
+# Corner-case tests for sum with infs and -infs.
+sum_inf
+
 # GPU backend uses floats to implement these ops for int32
 floor_int32
 divide_int32

src/ngraph/runtime/reference/sum.hpp

@@ -27,6 +27,29 @@ namespace ngraph
     {
         namespace reference
         {
+            // Windows doesn't seem to like it if we directly use std::isfinite on integer types,
+            // so we will roll our own thing here.
+            template <typename T>
+            typename std::enable_if<std::is_floating_point<T>::value, bool>::type is_finite(T x)
+            {
+                return std::isfinite(x);
+            }
+
+            template <typename T>
+            typename std::enable_if<std::is_same<T, bfloat16>::value ||
+                                        std::is_same<T, float16>::value,
+                                    bool>::type
+                is_finite(T x)
+            {
+                return std::isfinite(static_cast<float>(x));
+            }
+
+            template <typename T>
+            typename std::enable_if<std::is_integral<T>::value, bool>::type is_finite(T x)
+            {
+                return true;
+            }
+
             template <typename T>
             void sum(const T* arg,
                      T* out,
@@ -35,12 +58,12 @@ namespace ngraph
                      const AxisSet& reduction_axes)
             {
                 CoordinateTransform output_transform(out_shape);
-                std::vector<T> c(shape_size(out_shape));
+                std::vector<T> cs(shape_size(out_shape));
 
                 for (const Coordinate& output_coord : output_transform)
                 {
                     out[output_transform.index(output_coord)] = 0;
-                    c[output_transform.index(output_coord)] = 0;
+                    cs[output_transform.index(output_coord)] = 0;
                 }
 
                 CoordinateTransform input_transform(in_shape);
@@ -48,12 +71,21 @@ namespace ngraph
                 for (const Coordinate& input_coord : input_transform)
                 {
                     Coordinate output_coord = reduce(input_coord, reduction_axes);
-                    T y = arg[input_transform.index(input_coord)] -
-                          c[output_transform.index(output_coord)];
-                    T t = out[output_transform.index(output_coord)] + y;
-                    c[output_transform.index(output_coord)] =
-                        (t - out[output_transform.index(output_coord)]) - y;
-                    out[output_transform.index(output_coord)] = t;
+
+                    T x = arg[input_transform.index(input_coord)];
+                    T& z = out[output_transform.index(output_coord)];
+
+                    if (is_finite(x) && is_finite(z))
+                    {
+                        T& c = cs[output_transform.index(output_coord)];
+                        T t = z + (x - c);
+                        c = (t - z) - (x - c);
+                        z = t;
+                    }
+                    else
+                    {
+                        z = z + x;
+                    }
                 }
             }
         }

test/backend/sum.in.cpp

@@ -740,3 +740,39 @@ NGRAPH_TEST(${BACKEND_NAME}, sum_dynamic)
         ASSERT_TRUE(test::all_close_f(results, expected_results[i], MIN_FLOAT_TOLERANCE_BITS));
     }
 }
+
+NGRAPH_TEST(${BACKEND_NAME}, sum_inf)
+{
+    Shape shape{7, 4};
+    auto A = make_shared<op::Parameter>(element::f32, shape);
+    auto f = make_shared<Function>(make_shared<op::Sum>(A, AxisSet{1}), ParameterVector{A});
+
+    auto infi = std::numeric_limits<float>::infinity();
+
+    auto backend = runtime::Backend::create("${BACKEND_NAME}");
+
+    // Create some tensors for input/output
+    auto a = backend->create_tensor(element::f32, shape);
+    copy_data(a,
+              test::NDArray<float, 2>({{-infi, 0, 0, infi},
+                                       {infi, 100, -100, -infi},
+                                       {infi, 0, 100, infi},
+                                       {-infi, -100, 0, -infi},
+                                       {infi, infi, infi, infi},
+                                       {infi, infi, infi, -infi},
+                                       {infi, std::nanf(""), 42, infi}})
+                  .get_vector());
+    auto result = backend->create_tensor(element::f32, Shape{7});
+
+    auto handle = backend->compile(f);
+    handle->call_with_validate({result}, {a});
+    auto r = read_vector<float>(result);
+    ASSERT_EQ(r.size(), 7);
+    EXPECT_TRUE(isnan(r[0]));
+    EXPECT_TRUE(isnan(r[1]));
+    EXPECT_TRUE(r[2] > 0 && isinf(r[2]));
+    EXPECT_TRUE(r[3] < 0 && isinf(r[3]));
+    EXPECT_TRUE(r[4] > 0 && isinf(r[4]));
+    EXPECT_TRUE(isnan(r[5]));
+    EXPECT_TRUE(isnan(r[6]));
+}