std::tuple reduce_statistics()

in src/inplace_abn.cpp [24:56]

• 27 lines of code
• 3 McCabe index (conditional complexity)

std::tuple<at::Tensor, at::Tensor, at::Tensor> reduce_statistics(
    const at::Tensor& all_mean,
    const at::Tensor& all_var,
    const at::Tensor& all_count) {
  // Inputs shouldn't be half
  CHECK_NOT_HALF(all_mean);
  CHECK_NOT_HALF(all_var);

  // reduce_statistics is only used on GPU
  CHECK_CUDA(all_mean);
  CHECK_CUDA(all_var);
  CHECK_CUDA(all_count);

  // Check types and dimensions
  CHECK_SAME_TYPE(all_mean, all_var);
  IABN_CHECK(
      all_count.scalar_type() == at::ScalarType::Long,
      "all_count should have type int64");
  IABN_CHECK(all_mean.ndimension() == 2, "all_mean should have size N x C");
  IABN_CHECK(all_var.ndimension() == 2, "all_var should have size N x C");
  IABN_CHECK(
      all_count.ndimension() == 2 && all_count.size(1) == 1,
      "all_count should have size N x 1");
  IABN_CHECK(
      all_mean.size(0) == all_var.size(0) &&
          all_mean.size(0) == all_count.size(0),
      "Inputs should have the same size in dimension 0");
  IABN_CHECK(
      all_mean.size(1) == all_var.size(1),
      "all_mean and all_var should have the same size in dimension 1");

  return reduce_statistics_cuda(all_mean, all_var, all_count);
}