Function TestSampleMeanDistributed

tests/cpp/common/test_stats.cc:150–174

}

void TestSampleMeanDistributed(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;
      }
    }
    linalg::Vector<float> mean;
    SampleMean(&ctx, false, data, &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);
    }
  });
}

void TestWeightedSampleMean(Context const* ctx) {
  std::size_t m{32}, n{16};