LayerNormSecondPass Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/layer_norm_kernel.cpp lines 26–55

template <typename T>
void LayerNormSecondPass(
    const T* X_ptr,
    const T* gamma_data,
    const T* beta_data,
    T* Y_ptr,
    int64_t N,
    T scale,
    T bias) {
  using Vec = vec::Vectorized<T>;
  const bool gamma_null = gamma_data == nullptr;
  const bool beta_null = beta_data == nullptr;
  if (gamma_null || beta_null) {
    for (const auto j : c10::irange(N)) {
      const T gamma_v = gamma_null ? T(1) : gamma_data[j];
      const T beta_v = beta_null ? T(0) : beta_data[j];
      Y_ptr[j] = (X_ptr[j] + bias) * scale * gamma_v + beta_v;
    }
  } else {
    vec::map3<T>(
      [scale, bias](Vec x, Vec gamma, Vec beta) {
        return (x + Vec(bias)) * Vec(scale) * gamma + beta;
      },
      Y_ptr,
      X_ptr,
      gamma_data,
      beta_data,
      N);
  }
}