class_docstr: ## Batch Normalization Layer
  
  Batch normalization layer $\text{BN}$ normalizes the input $X$ as follows:
  
  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.
  $\gamma \in \mathbb{R}^{C}$ and $\beta \in \mathbb{R}^{C}$.
  $$\text{BN}(X) = \gamma
  \frac{X - \underset{B, H, W}{\mathbb{E}}[X]}{\sqrt{\underset{B, H, W}{Var}[X] + \epsilon}}
  + \beta$$
  
  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{BN}(X) = \gamma
  \frac{X - \underset{B}{\mathbb{E}}[X]}{\sqrt{\underset{B}{Var}[X] + \epsilon}}
  + \beta$$
  
  When input $X \in \mathbb{R}^{B \times C \times L}$ is a batch of a sequence embeddings,
  where $B$ is the batch size, $C$ is the number of features, and $L$ is the length of the sequence.
  $\gamma \in \mathbb{R}^{C}$ and $\beta \in \mathbb{R}^{C}$.
  $$\text{BN}(X) = \gamma
  \frac{X - \underset{B, L}{\mathbb{E}}[X]}{\sqrt{\underset{B, L}{Var}[X] + \epsilon}}
  + \beta$$
docstr: `x` is a tensor of shape `[batch_size, channels, *]`.
  `*` denotes any number of (possibly 0) dimensions.
   For example, in an image (2D) convolution this will be
  `[batch_size, channels, height, width]`
docstr_tok: ["`","x","`","is","a","tensor","of","shape","`","[","batch_size",",","channels",",","*","]","`",".","`","*","`","denotes","any","number","of","(","possibly","0",")","dimensions",".","For","example",",","in","an","image","(","2D",")","convolution","this","will","be","`","[","batch_size",",","channels",",","height",",","width","]","`"]
code: def forward(self, x: torch.Tensor):
      """
      `x` is a tensor of shape `[batch_size, channels, *]`.
      `*` denotes any number of (possibly 0) dimensions.
       For example, in an image (2D) convolution this will be
      `[batch_size, channels, height, width]`
      """
      # Keep the original shape
      x_shape = x.shape
      # Get the batch size
      batch_size = x_shape[0]
      # Sanity check to make sure the number of features is the same
      assert self.channels == x.shape[1]
  
      # Reshape into `[batch_size, channels, n]`
      x = x.view(batch_size, self.channels, -1)
  
      # We will calculate the mini-batch mean and variance
      # if we are in training mode or if we have not tracked exponential moving averages
      if self.training or not self.track_running_stats:
          # Calculate the mean across first and last dimension;
          # i.e. the means for each feature $\mathbb{E}[x^{(k)}]$
          mean = x.mean(dim=[0, 2])
          # Calculate the squared mean across first and last dimension;
          # i.e. the means for each feature $\mathbb{E}[(x^{(k)})^2]$
          mean_x2 = (x**2).mean(dim=[0, 2])
          # Variance for each feature $Var[x^{(k)}] = \mathbb{E}[(x^{(k)})^2] - \mathbb{E}[x^{(k)}]^2$
          var = mean_x2 - mean**2
  
          # Update exponential moving averages
          if self.training and self.track_running_stats:
              self.exp_mean = (
                  1 - self.momentum
              ) * self.exp_mean + self.momentum * mean
              self.exp_var = (
                  1 - self.momentum
              ) * self.exp_var + self.momentum * var
      # Use exponential moving averages as estimates
      else:
          mean = self.exp_mean
          var = self.exp_var
  
      # Normalize $$\hat{x}^{(k)} = \frac{x^{(k)} - \mathbb{E}[x^{(k)}]}{\sqrt{Var[x^{(k)}] + \epsilon}}$$
      x_norm = (x - mean.view(1, -1, 1)) / torch.sqrt(
          var + self.eps
      ).view(1, -1, 1)
      # Scale and shift $$y^{(k)} =\gamma^{(k)} \hat{x}^{(k)} + \beta^{(k)}$$
      if self.affine:
          x_norm = self.scale.view(1, -1, 1) * x_norm + self.shift.view(
              1, -1, 1
          )
  
      # Reshape to original and return
      return x_norm.view(x_shape)
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/normalization/batch_norm/__init__.py
imports: import torch
  from torch import nn
  from labml_helpers.module import Module
label_desc: Use image node_id 2 for calling the BatchNorm obj's underlying member method code with example usage: obj.forward(x) and returns: x_norm