Function TestIterationRangeColumnSplit

tests/cpp/predictor/test_predictor.cc:579–630

}  // anonymous namespace

void TestIterationRangeColumnSplit(int world_size, bool use_gpu) {
  std::size_t constexpr kRows = 1000, kCols = 20, kClasses = 4, kForest = 3, kIters = 10;
  auto dmat = RandomDataGenerator(kRows, kCols, 0).Classes(kClasses).GenerateDMatrix(true);
  Context ctx;
  if (use_gpu) {
    ctx = MakeCUDACtx(0);
  }
  auto learner = LearnerForTest(&ctx, dmat, kIters, kForest);

  bool bound = false;
  std::unique_ptr<Learner> sliced{learner->Slice(0, 3, 1, &bound)};
  ASSERT_FALSE(bound);

  // margin
  HostDeviceVector<float> margin_predt_sliced;
  HostDeviceVector<float> margin_predt_ranged;
  sliced->Predict(dmat, true, &margin_predt_sliced, 0, 0, false, false, false, false, false);
  learner->Predict(dmat, true, &margin_predt_ranged, 0, 3, false, false, false, false, false);
  auto const &margin_sliced = margin_predt_sliced.HostVector();
  auto const &margin_ranged = margin_predt_ranged.HostVector();

  // Leaf
  HostDeviceVector<float> leaf_predt_sliced;
  HostDeviceVector<float> leaf_predt_ranged;
  sliced->Predict(dmat, false, &leaf_predt_sliced, 0, 0, false, true, false, false, false);
  learner->Predict(dmat, false, &leaf_predt_ranged, 0, 3, false, true, false, false, false);
  auto const &leaf_sliced = leaf_predt_sliced.HostVector();
  auto const &leaf_ranged = leaf_predt_ranged.HostVector();

  Json ranged_model{Object{}};
  learner->SaveModel(&ranged_model);
  Json sliced_model{Object{}};
  sliced->SaveModel(&sliced_model);

#if !defined(XGBOOST_USE_NCCL)
  if (use_gpu) {
    GTEST_SKIP_("Not compiled with NCCL");
    return;
  }
#endif  // defined(XGBOOST_USE_NCCL)
  collective::TestDistributedGlobal(world_size, [&] {
    VerifyIterationRangeColumnSplit(use_gpu, ranged_model, sliced_model, kRows, kCols, kClasses,
                                    margin_ranged, margin_sliced, leaf_ranged, leaf_sliced);
  });

#if defined(XGBOOST_USE_FEDERATED)
  collective::TestFederatedGlobal(world_size, [&] {
    VerifyIterationRangeColumnSplit(use_gpu, ranged_model, sliced_model, kRows, kCols, kClasses,
                                    margin_ranged, margin_sliced, leaf_ranged, leaf_sliced);
  });
#endif  // defined(XGBOOST_USE_FEDERATED)
}

void TestSparsePrediction(Context const *ctx, float sparsity) {
  size_t constexpr kRows = 512, kCols = 128, kIters = 4;