add results rows checks + early return

ginkgo-project · Jun 27, 2024 · 872b7ec · 872b7ec
1 parent 1d813ee
commit 872b7ec
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/reference/matrix/dense_kernels.cpp b/reference/matrix/dense_kernels.cpp
@@ -374,11 +374,15 @@ void compute_mean(std::shared_ptr<const ReferenceExecutor> exec,
                   const matrix::Dense<ValueType>* x,
                   matrix::Dense<ValueType>* result, array<char>&)
 {
+    GKO_ASSERT_EQ(result->get_size()[0], 1);
+
     using ValueType_nc = gko::remove_complex<ValueType>;
     for (size_type j = 0; j < x->get_size()[1]; ++j) {
         result->at(0, j) = zero<ValueType>();
     }
 
+    if (x->get_size()[0] == 0) return;
+
     for (size_type i = 0; i < x->get_size()[1]; ++i) {
         for (size_type j = 0; j < x->get_size()[0]; ++j) {
             result->at(0, i) += x->at(j, i);

diff --git a/reference/test/matrix/dense_kernels.cpp b/reference/test/matrix/dense_kernels.cpp
@@ -695,6 +695,14 @@ TYPED_TEST(Dense, ComputesMean)
     GKO_EXPECT_NEAR(result->at(0, 2), T{1.0}, r<T>::value * 10);
 }
 
+TYPED_TEST(Dense, ComputesMeanFailsOnZeroRowResults)
+{
+    using Mtx = typename TestFixture::Mtx;
+    using T = typename TestFixture::value_type;
+    auto result = Mtx::create(this->exec, gko::dim<2>{0, 1});
+
+    ASSERT_THROW(this->mtx4->compute_mean(result), gko::ValueMismatch);
+}
 
 TYPED_TEST(Dense, ComputesMeanFailsOnWrongResultSize)
 {