CalcMeanVar Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/group_norm_kernel.cpp lines 162–182

template <typename T, typename opmath_t>
inline std::enable_if_t<std::is_same_v<T, opmath_t>, void>
CalcMeanVar(
  const T* X_ptr,
  opmath_t* mean_ptr,
  opmath_t* rstd_ptr,
  int64_t C) {
  using Vec = vec::Vectorized<T>;
  vec::map2<T>(
          [](Vec x, Vec y) { return x + y; },
          mean_ptr,
          X_ptr,
          mean_ptr,
          C);
  vec::map2<T>(
      [](Vec x, Vec y) { return x * x + y; },
      rstd_ptr,
      X_ptr,
      rstd_ptr,
      C);
}