scalar_inner_sum Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/SumKernel.cpp lines 463–473

template <typename acc_t, typename LoadPolicy, typename StorePolicy>
void scalar_inner_sum(
    // NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays)
    char * C10_RESTRICT data[2], int64_t in_strides[2], int64_t out_stride,
    int64_t size0, int64_t size1) {
  for (const auto j : c10::irange(size1)) {
    const auto *row_in = data[1] + j * in_strides[1];
    auto ans = row_sum<acc_t, LoadPolicy>(row_in, in_strides[0], size0);
    store<StorePolicy>(data[0], out_stride, j, ans);
  }
}