ApplyScaleBias Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/group_norm_kernel.cpp lines 229–245

template <typename T, typename opmath_t>
inline std::enable_if_t<std::is_same_v<T, opmath_t>, void>
ApplyScaleBias(
  T* Y_ptr,
  const T* X_ptr,
  const opmath_t* scale_ptr,
  const opmath_t* bias_ptr,
  int64_t C) {
  using Vec = vec::Vectorized<T>;
  vec::map3<T>(
    [](Vec x, Vec scale, Vec bias) { return x * scale + bias; },
    Y_ptr,
    X_ptr,
    scale_ptr,
    bias_ptr,
    C);
}