class_docstr: ## Channel Normalization
  
  This is similar to [Group Normalization](../group_norm/index.html) but affine transform is done group wise.
class_docstr_tok: ["#","#","Channel","Normalization","This","is","similar","to","[","Group","Normalization","]","(","..","\/group_norm\/index.html",")","but","affine","transform","is","done","group","wise","."]
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, groups, n]`
      x = x.view(batch_size, self.groups, -1)
  
      # Calculate the mean across last dimension;
      # i.e. the means for each sample and channel group $\mathbb{E}[x_{(i_N, i_G)}]$
      mean = x.mean(dim=[-1], keepdim=True)
      # Calculate the squared mean across last dimension;
      # i.e. the means for each sample and channel group $\mathbb{E}[x^2_{(i_N, i_G)}]$
      mean_x2 = (x**2).mean(dim=[-1], keepdim=True)
      # Variance for each sample and feature group
      # $Var[x_{(i_N, i_G)}] = \mathbb{E}[x^2_{(i_N, i_G)}] - \mathbb{E}[x_{(i_N, i_G)}]^2$
      var = mean_x2 - mean**2
  
      # Normalize
      # $$\hat{x}_{(i_N, i_G)} =
      # \frac{x_{(i_N, i_G)} - \mathbb{E}[x_{(i_N, i_G)}]}{\sqrt{Var[x_{(i_N, i_G)}] + \epsilon}}$$
      x_norm = (x - mean) / torch.sqrt(var + self.eps)
  
      # Scale and shift group-wise
      # $$y_{i_G} =\gamma_{i_G} \hat{x}_{i_G} + \beta_{i_G}$$
      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 ChannelNorm obj's underlying member method code with example usage: obj.forward(x) and returns: x_norm