class_docstr: ## Layer Normalization
  
  Layer normalization $\text{LN}$ normalizes the input $X$ as follows:
  
  When input $X \in \mathbb{R}^{B \times C}$ is a batch of embeddings,
  where $B$ is the batch size and $C$ is the number of features.
  $\gamma \in \mathbb{R}^{C}$ and $\beta \in \mathbb{R}^{C}$.
  $$\text{LN}(X) = \gamma
  \frac{X - \underset{C}{\mathbb{E}}[X]}{\sqrt{\underset{C}{Var}[X] + \epsilon}}
  + \beta$$
  
  When input $X \in \mathbb{R}^{L \times B \times C}$ is a batch of a sequence of embeddings,
  where $B$ is the batch size, $C$ is the number of channels, $L$ is the length of the sequence.
  $\gamma \in \mathbb{R}^{C}$ and $\beta \in \mathbb{R}^{C}$.
  $$\text{LN}(X) = \gamma
  \frac{X - \underset{C}{\mathbb{E}}[X]}{\sqrt{\underset{C}{Var}[X] + \epsilon}}
  + \beta$$
  
  When input $X \in \mathbb{R}^{B \times C \times H \times W}$ is a batch of image representations,
  where $B$ is the batch size, $C$ is the number of channels, $H$ is the height and $W$ is the width.
  This is not a widely used scenario.
  $\gamma \in \mathbb{R}^{C \times H \times W}$ and $\beta \in \mathbb{R}^{C \times H \times W}$.
  $$\text{LN}(X) = \gamma
  \frac{X - \underset{C, H, W}{\mathbb{E}}[X]}{\sqrt{\underset{C, H, W}{Var}[X] + \epsilon}}
  + \beta$$
docstr: `x` is a tensor of shape `[*, S[0], S[1], ..., S[n]]`.
  `*` could be any number of dimensions.
   For example, in an NLP task this will be
  `[seq_len, batch_size, features]`
code: def forward(self, x: torch.Tensor):
      """
      `x` is a tensor of shape `[*, S[0], S[1], ..., S[n]]`.
      `*` could be any number of dimensions.
       For example, in an NLP task this will be
      `[seq_len, batch_size, features]`
      """
      # Sanity check to make sure the shapes match
      assert (
          self.normalized_shape
          == x.shape[-len(self.normalized_shape) :]
      )
  
      # The dimensions to calculate the mean and variance on
      dims = [-(i + 1) for i in range(len(self.normalized_shape))]
  
      # Calculate the mean of all elements;
      # i.e. the means for each element $\mathbb{E}[X]$
      mean = x.mean(dim=dims, keepdim=True)
      # Calculate the squared mean of all elements;
      # i.e. the means for each element $\mathbb{E}[X^2]$
      mean_x2 = (x**2).mean(dim=dims, keepdim=True)
      # Variance of all element $Var[X] = \mathbb{E}[X^2] - \mathbb{E}[X]^2$
      var = mean_x2 - mean**2
  
      # Normalize $$\hat{X} = \frac{X - \mathbb{E}[X]}{\sqrt{Var[X] + \epsilon}}$$
      x_norm = (x - mean) / torch.sqrt(var + self.eps)
      # Scale and shift $$\text{LN}(x) = \gamma \hat{X} + \beta$$
      if self.elementwise_affine:
          x_norm = self.gain * x_norm + self.bias
  
      #
      return x_norm
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/normalization/layer_norm/__init__.py
imports: from typing import Union, List
  import torch
  from torch import nn, Size
  from labml_helpers.module import Module
label_desc: Use image node_id 2 for calling the LayerNorm obj's underlying member method code with example usage: obj.forward(x) and returns: x_norm