Function TestWeightedSampleMeanDistributed

tests/cpp/common/test_stats.cc:203–229

}

void TestWeightedSampleMeanDistributed(Context const* ctx) {
  std::size_t m{32}, n{16};
  auto device = ctx->Device();
  std::int32_t n_workers =
      device.IsCPU() ? std::min(4u, std::thread::hardware_concurrency()) : curt::AllVisibleGPUs();

  collective::TestDistributedGlobal(n_workers, [m, n, device, n_workers] {
    auto rank = collective::GetRank();
    Context ctx = device.IsCUDA() ? MakeCUDACtx(DistGpuIdx()) : Context{};
    collective::GetWorkerLocalThreads(collective::GetWorldSize(), &ctx);
    linalg::Matrix<float> data({m, n}, ctx.Device());
    auto h_data = data.HostView();
    for (std::size_t i = 0; i < m; ++i) {
      for (std::size_t j = 0; j < n; ++j) {
        h_data(i, j) = i + (m * rank) + j;
      }
    }
    HostDeviceVector<float> w{m, 1.0f, ctx.Device()};
    linalg::Vector<float> mean;
    WeightedSampleMean(&ctx, false, data, w, &mean);
    ASSERT_EQ(mean.Size(), n);
    double total = n_workers * m;
    for (std::size_t i = 0; i < n; ++i) {
      ASSERT_EQ(mean(i), (i + total - 1.0 + i) * total / 2.0 / total);
    }
  });
}
}  // namespace

TEST(Stats, SampleMean) {