class_docstr: ## Estimated Batch Normalization
  
  When input $X \in \mathbb{R}^{B     imes C  imes H  imes 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 $eta \in \mathbb{R}^{C}$.
  
  $$\dot{X}_{\cdot, C, \cdot, \cdot} = \gamma_C
  rac{X_{\cdot, C, \cdot, \cdot} - \hat{\mu}_C}{\hat{\sigma}_C}
  + eta_C$$
  
  where,
  
  egin{align}
  \hat{\mu}_C &\longleftarrow (1 - r)\hat{\mu}_C + r rac{1}{B H W} \sum_{b,h,w} X_{b,c,h,w} \
  \hat{\sigma}^2_C &\longleftarrow (1 - r)\hat{\sigma}^2_C + r rac{1}{B H W} \sum_{b,h,w} ig(X_{b,c,h,w} - \hat{\mu}_C ig)^2
  \end{align}
  
  are the running mean and variances. $r$ is the momentum for calculating the exponential mean.
class_docstr_tok: ["#","#","Estimated","Batch","Normalization","When","input","$","X","\\in","\\mathbb","{","R","}","^","{","B","imes","C","imes","H","imes","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","}","$",".","$","$","\\dot","{","X","}","_","{","\\cdot",",","C",",","\\cdot",",","\\cdot","}","=","\\gamma_C","rac","{","X_","{","\\cdot",",","C",",","\\cdot",",","\\cdot","}","-","\\hat","{","\\mu","}","_C","}","{","\\hat","{","\\sigma","}","_C","}","+","\beta_C","$","$","where",",","\begin","{","align","}","\\hat","{","\\mu","}","_C","&","\\longleftarrow","(","1","-","r",")","\\hat","{","\\mu","}","_C","+","r","rac","{","1","}","{","B","H","W","}","\\sum_","{","b",",","h",",","w","}","X_","{","b",",","c",",","h",",","w","}","\\","\\hat","{","\\sigma","}","^2_C","&","\\longleftarrow","(","1","-","r",")","\\hat","{","\\sigma","}","^2_C","+","r","rac","{","1","}","{","B","H","W","}","\\sum_","{","b",",","h",",","w","}","\big","(","X_","{","b",",","c",",","h",",","w","}","-","\\hat","{","\\mu","}","_C","\big",")","^2","\\end","{","align","}","are","the","running","mean","and","variances",".","$","r","$","is","the","momentum","for","calculating","the","exponential","mean","."]
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 old shape
      x_shape = x.shape
      # Get the batch size
      batch_size = x_shape[0]
  
      # Sanity check to make sure the number of features is correct
      assert self.channels == x.shape[1]
  
      # Reshape into `[batch_size, channels, n]`
      x = x.view(batch_size, self.channels, -1)
  
      # Update $\hat{\mu}_C$ and $\hat{\sigma}^2_C$ in training mode only
      if self.training:
          # No backpropagation through $\hat{\mu}_C$ and $\hat{\sigma}^2_C$
          with torch.no_grad():
              # Calculate the mean across first and last dimensions;
              # $$\frac{1}{B H W} \sum_{b,h,w} X_{b,c,h,w}$$
              mean = x.mean(dim=[0, 2])
              # Calculate the squared mean across first and last dimensions;
              # $$\frac{1}{B H W} \sum_{b,h,w} X^2_{b,c,h,w}$$
              mean_x2 = (x**2).mean(dim=[0, 2])
              # Variance for each feature
              # $$\frac{1}{B H W} \sum_{b,h,w} \big(X_{b,c,h,w} - \hat{\mu}_C \big)^2$$
              var = mean_x2 - mean**2
  
              # Update exponential moving averages
              #
              # \begin{align}
              # \hat{\mu}_C &\longleftarrow (1 - r)\hat{\mu}_C + r \frac{1}{B H W} \sum_{b,h,w} X_{b,c,h,w} \\
              # \hat{\sigma}^2_C &\longleftarrow (1 - r)\hat{\sigma}^2_C + r \frac{1}{B H W} \sum_{b,h,w} \big(X_{b,c,h,w} - \hat{\mu}_C \big)^2
              # \end{align}
              self.exp_mean = (
                  1 - self.momentum
              ) * self.exp_mean + self.momentum * mean
              self.exp_var = (
                  1 - self.momentum
              ) * self.exp_var + self.momentum * var
  
      # Normalize
      # $$\frac{X_{\cdot, C, \cdot, \cdot} - \hat{\mu}_C}{\hat{\sigma}_C}$$
      x_norm = (x - self.exp_mean.view(1, -1, 1)) / torch.sqrt(
          self.exp_var + self.eps
      ).view(1, -1, 1)
      # Scale and shift
      # $$ \gamma_C
      #     \frac{X_{\cdot, C, \cdot, \cdot} - \hat{\mu}_C}{\hat{\sigma}_C}
      #     + \beta_C$$
      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_channel_norm/__init__.py
imports: import torch
  from torch import nn
  from labml_helpers.module import Module
  from labml_nn.normalization.batch_norm import BatchNorm
label_desc: Use image node_id 2 for calling the EstimatedBatchNorm obj's underlying member method code with example usage: obj.forward(x) and returns: x_norm