ComputeInternalGradients Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/group_norm_kernel.cpp lines 531–556

template <typename T, typename opmath_t>
std::enable_if_t<std::is_same_v<T, opmath_t>, void>
ComputeInternalGradients(
    int64_t N,
    int64_t C,
    int64_t HxW,
    const T* dY,
    const T* X,
    opmath_t* ds,
    opmath_t* db) {
  using Vec = at::vec::Vectorized<opmath_t>;
  at::parallel_for(0, N * C, 1, [=](int64_t start, int64_t end) {
    for (const auto i : c10::irange(start, end)) {
      const T* dY_ptr = dY + i * HxW;
      const T* X_ptr = X + i * HxW;
      ds[i] = at::vec::map2_reduce_all<T>(
          [](Vec x, Vec y) { return x * y; },
          [](Vec x, Vec y) { return x + y; },
          dY_ptr,
          X_ptr,
          HxW);
      db[i] = at::vec::reduce_all<T>(
          [](Vec& x, Vec& y) { return x + y; }, dY_ptr, HxW);
    }
  });
}