LayerNormKernelImplInternal Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/layer_norm_kernel.cpp lines 156–174

template <typename T,
          typename std::enable_if_t<is_reduced_floating_point_v<T>, int> = 0>
void LayerNormKernelImplInternal(
    const Tensor& X,
    const Tensor& gamma,
    const Tensor& beta,
    int64_t M,
    int64_t N,
    float eps,
    Tensor* Y,
    Tensor* mean,
    Tensor* rstd) {
  const bool mixed_type = is_mixed_type(X, gamma, beta);
  if (mixed_type) {
    layer_norm_kernel_mixed_type<T, float>(X, gamma, beta, M, N, eps, Y, mean, rstd);
  } else {
    layer_norm_kernel_mixed_type<T, T>(X, gamma, beta, M, N, eps, Y, mean, rstd);
  }
}