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: * `channels` is the number of features in the input
  * `eps` is $\epsilon$, used in $\sqrt{Var[x^{(k)}] + \epsilon}$ for numerical stability
  * `momentum` is the momentum in taking the exponential moving average
  * `affine` is whether to scale and shift the normalized value
  * `track_running_stats` is whether to calculate the moving averages or mean and variance
  
  We've tried to use the same names for arguments as PyTorch `BatchNorm` implementation.
docstr_tok: ["*","`","channels","`","is","the","number","of","features","in","the","input","*","`","eps","`","is","$","\\epsilon","$",",","used","in","$","\\sqrt","{","Var","[","x^","{","(","k",")","}","]","+","\\epsilon","}","$","for","numerical","stability","*","`","momentum","`","is","the","momentum","in","taking","the","exponential","moving","average","*","`","affine","`","is","whether","to","scale","and","shift","the","normalized","value","*","`","track_running_stats","`","is","whether","to","calculate","the","moving","averages","or","mean","and","variance","We","'ve","tried","to","use","the","same","names","for","arguments","as","PyTorch","`","BatchNorm","`","implementation","."]
code: def __init__(
      self,
      channels: int,
      *,
      eps: float = 1e-5,
      momentum: float = 0.1,
      affine: bool = True,
      track_running_stats: bool = True,
  ):
      """
      * `channels` is the number of features in the input
      * `eps` is $\epsilon$, used in $\sqrt{Var[x^{(k)}] + \epsilon}$ for numerical stability
      * `momentum` is the momentum in taking the exponential moving average
      * `affine` is whether to scale and shift the normalized value
      * `track_running_stats` is whether to calculate the moving averages or mean and variance
  
      We've tried to use the same names for arguments as PyTorch `BatchNorm` implementation.
      """
      super().__init__()
  
      self.channels = channels
  
      self.eps = eps
      self.momentum = momentum
      self.affine = affine
      self.track_running_stats = track_running_stats
      # Create parameters for $\gamma$ and $\beta$ for scale and shift
      if self.affine:
          self.scale = nn.Parameter(torch.ones(channels))
          self.shift = nn.Parameter(torch.zeros(channels))
      # Create buffers to store exponential moving averages of
      # mean $\mathbb{E}[x^{(k)}]$ and variance $Var[x^{(k)}]$
      if self.track_running_stats:
          self.register_buffer("exp_mean", torch.zeros(channels))
          self.register_buffer("exp_var", torch.ones(channels))
  
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 1 to create a new BatchNorm object from inherited base classes: Module with example: obj = BatchNorm(channels)