Function VerifyIterationRangeColumnSplit

tests/cpp/predictor/test_predictor.cc:515–576

namespace {
void VerifyIterationRangeColumnSplit(bool use_gpu, Json const &ranged_model,
                                     Json const &sliced_model, std::size_t rows, std::size_t cols,
                                     std::size_t classes,
                                     std::vector<float> const &expected_margin_ranged,
                                     std::vector<float> const &expected_margin_sliced,
                                     std::vector<float> const &expected_leaf_ranged,
                                     std::vector<float> const &expected_leaf_sliced) {
  auto const world_size = collective::GetWorldSize();
  auto const rank = collective::GetRank();
  Context ctx;
  if (use_gpu) {
    ctx = MakeCUDACtx(curt::AllVisibleGPUs() == 1 ? 0 : rank);
  }
  collective::GetWorkerLocalThreads(world_size, &ctx);

  auto dmat = RandomDataGenerator(rows, cols, 0).Classes(classes).GenerateDMatrix(true);
  std::shared_ptr<DMatrix> Xy{dmat->SliceCol(world_size, rank)};

  std::unique_ptr<Learner> learner{Learner::Create({Xy})};
  auto args = Args{{"device", ctx.DeviceName()}, {"nthread", std::to_string(ctx.Threads())}};
  learner->SetParams(args);
  learner->LoadModel(ranged_model);

  std::unique_ptr<Learner> sliced{Learner::Create({Xy})};
  sliced->SetParams(args);
  sliced->LoadModel(sliced_model);

  HostDeviceVector<float> out_predt_sliced;
  HostDeviceVector<float> out_predt_ranged;

  // margin
  {
    sliced->Predict(Xy, true, &out_predt_sliced, 0, 0, false, false, false, false, false);
    learner->Predict(Xy, true, &out_predt_ranged, 0, 3, false, false, false, false, false);
    auto const &h_sliced = out_predt_sliced.HostVector();
    auto const &h_ranged = out_predt_ranged.HostVector();
    EXPECT_EQ(h_sliced.size(), expected_margin_sliced.size());
    for (std::size_t i = 0; i < expected_margin_sliced.size(); ++i) {
      ASSERT_FLOAT_EQ(h_sliced[i], expected_margin_sliced[i]) << "rank " << rank << ", i " << i;
    }
    EXPECT_EQ(h_ranged.size(), expected_margin_ranged.size());
    for (std::size_t i = 0; i < expected_margin_ranged.size(); ++i) {
      ASSERT_FLOAT_EQ(h_ranged[i], expected_margin_ranged[i]) << "rank " << rank << ", i " << i;
    }
  }

  // Leaf
  {
    sliced->Predict(Xy, false, &out_predt_sliced, 0, 0, false, true, false, false, false);
    learner->Predict(Xy, false, &out_predt_ranged, 0, 3, false, true, false, false, false);
    auto const &h_sliced = out_predt_sliced.HostVector();
    auto const &h_ranged = out_predt_ranged.HostVector();
    EXPECT_EQ(h_sliced.size(), expected_leaf_sliced.size());
    for (std::size_t i = 0; i < expected_leaf_sliced.size(); ++i) {
      ASSERT_FLOAT_EQ(h_sliced[i], expected_leaf_sliced[i]) << "rank " << rank << ", i " << i;
    }
    EXPECT_EQ(h_ranged.size(), expected_leaf_ranged.size());
    for (std::size_t i = 0; i < expected_leaf_ranged.size(); ++i) {