Use numerically stable tree reduction in TensorReduction.

1 files changed, 24 insertions, 4 deletions

diff --git a/unsupported/test/cxx11_tensor_reduction.cpp b/unsupported/test/cxx11_tensor_reduction.cpp
index ff8e18c07..9458f4e4e 100644
--- a/unsupported/test/cxx11_tensor_reduction.cpp
+++ b/unsupported/test/cxx11_tensor_reduction.cpp
@@ -386,7 +386,7 @@ static void test_static_dims() {
           expected = (std::max)(expected, in(i, k, j, l));
         }
       }
-      VERIFY_IS_APPROX(out(i, j), expected);
+      VERIFY_IS_EQUAL(out(i, j), expected);
     }
   }
 }
@@ -417,7 +417,7 @@ static void test_innermost_last_dims() {
           expected = (std::max)(expected, in(l, k, i, j));
         }
       }
-      VERIFY_IS_APPROX(out(i, j), expected);
+      VERIFY_IS_EQUAL(out(i, j), expected);
     }
   }
 }
@@ -448,7 +448,7 @@ static void test_innermost_first_dims() {
           expected = (std::max)(expected, in(i, j, k, l));
         }
       }
-      VERIFY_IS_APPROX(out(i, j), expected);
+      VERIFY_IS_EQUAL(out(i, j), expected);
     }
   }
 }
@@ -479,11 +479,30 @@ static void test_reduce_middle_dims() {
           expected = (std::max)(expected, in(i, k, l, j));
         }
       }
-      VERIFY_IS_APPROX(out(i, j), expected);
+      VERIFY_IS_EQUAL(out(i, j), expected);
     }
   }
 }
 
+static void test_sum_accuracy() {
+  Tensor<float, 3> tensor(101, 101, 101);
+  for (float prescribed_mean : {1.0f, 10.0f, 100.0f, 1000.0f, 10000.0f}) {
+    tensor.setRandom();
+    tensor += tensor.constant(prescribed_mean);
+
+    Tensor<float, 0> sum = tensor.sum();
+    double expected_sum = 0.0;
+    for (int i = 0; i < 101; ++i) {
+      for (int j = 0; j < 101; ++j) {
+        for (int k = 0; k < 101; ++k) {
+          expected_sum += static_cast<double>(tensor(i, j, k));
+        }
+      }
+    }
+    VERIFY_IS_APPROX(sum(), static_cast<float>(expected_sum));
+  }
+}
+
 EIGEN_DECLARE_TEST(cxx11_tensor_reduction) {
   CALL_SUBTEST(test_trivial_reductions<ColMajor>());
   CALL_SUBTEST(test_trivial_reductions<RowMajor>());
@@ -506,4 +525,5 @@ EIGEN_DECLARE_TEST(cxx11_tensor_reduction) {
   CALL_SUBTEST(test_innermost_first_dims<RowMajor>());
   CALL_SUBTEST(test_reduce_middle_dims<ColMajor>());
   CALL_SUBTEST(test_reduce_middle_dims<RowMajor>());
+  CALL_SUBTEST(test_sum_accuracy());
 }