From 9794ec6b8e38a176e37cbf95d1355edbef71e5e9 Mon Sep 17 00:00:00 2001
From: Joona Havukainen <joona.j.havukainen@gmail.com>
Date: Sun, 6 Jul 2025 16:04:50 -0700
Subject: [PATCH] Improve the cpp unittest

---
 tests/ops_tests.cpp | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp
index af3e3a698..1a9781c7c 100644
--- a/tests/ops_tests.cpp
+++ b/tests/ops_tests.cpp
@@ -1025,8 +1025,9 @@ TEST_CASE("test reduction ops") {
     CHECK(array_equal(min(x, 1), array({true, false})).item<bool>());
     CHECK(array_equal(min(x, 0), array({false, true, false})).item<bool>());
 
-    x = array({1.0f, NAN, 3.0f, 4.0f, 5.0f, 6.0f});
-    CHECK(isnan(max(x).item<float>()));
+    x = array({1.0f, NAN, 3.0f, 4.0f, 5.0f, 6.0f}, {2, 3});
+    CHECK(array_equal(max(x, 0), array({4.0f, NAN, 6.0f}), true).item<bool>());
+    CHECK(array_equal(max(x, 1), array({NAN, 6.0f}), true).item<bool>());
   }
 
   // Test logsumexp