merge latest, Jan 12 2024

tests/cpp/common/test_device_helpers.cu

@@ -171,7 +171,7 @@ TEST(SegmentedUnique, Regression) {
   }
 }
 
-TEST(Allocator, OOM) {
+TEST(Allocator, DISABLED_OOM) {
   auto size = dh::AvailableMemory(0) * 4;
   ASSERT_THROW({dh::caching_device_vector<char> vec(size);}, dmlc::Error);
   ASSERT_THROW({dh::device_vector<char> vec(size);}, dmlc::Error);

tests/cpp/common/test_linalg.cu

@@ -22,7 +22,7 @@ void TestElementWiseKernel() {
     ElementWiseTransformDevice(t, [] __device__(size_t i, float) { return i; });
     // CPU view
     t = l.View(DeviceOrd::CPU()).Slice(linalg::All(), 1, linalg::All());
-    size_t k = 0;
+    std::size_t k = 0;
     for (size_t i = 0; i < l.Shape(0); ++i) {
       for (size_t j = 0; j < l.Shape(2); ++j) {
         ASSERT_EQ(k++, t(i, j));
@@ -30,7 +30,15 @@ void TestElementWiseKernel() {
     }
 
     t = l.View(device).Slice(linalg::All(), 1, linalg::All());
-    ElementWiseKernelDevice(t, [] XGBOOST_DEVICE(size_t i, float v) { SPAN_CHECK(v == i); });
+    cuda_impl::ElementWiseKernel(
+        t, [=] XGBOOST_DEVICE(std::size_t i, std::size_t j) mutable { t(i, j) = i + j; });
+
+    t = l.Slice(linalg::All(), 1, linalg::All());
+    for (size_t i = 0; i < l.Shape(0); ++i) {
+      for (size_t j = 0; j < l.Shape(2); ++j) {
+        ASSERT_EQ(i + j, t(i, j));
+      }
+    }
   }
 
   {