class_docstr: Applies batch normalization over a batch of heterogeneous features as
  described in the `"Batch Normalization: Accelerating Deep Network Training
  by Reducing Internal Covariate Shift" <https://arxiv.org/abs/1502.03167>`_
  paper.
  Compared to :class:`BatchNorm`, :class:`HeteroBatchNorm` applies
  normalization individually for each node or edge type.
  
  Args:
      in_channels (int): Size of each input sample.
      num_types (int): The number of types.
      eps (float, optional): A value added to the denominator for numerical
          stability. (default: :obj:`1e-5`)
      momentum (float, optional): The value used for the running mean and
          running variance computation. (default: :obj:`0.1`)
      affine (bool, optional): If set to :obj:`True`, this module has
          learnable affine parameters :math:`\gamma` and :math:`\beta`.
          (default: :obj:`True`)
      track_running_stats (bool, optional): If set to :obj:`True`, this
          module tracks the running mean and variance, and when set to
          :obj:`False`, this module does not track such statistics and always
          uses batch statistics in both training and eval modes.
          (default: :obj:`True`)
class_docstr_tok: ["Applies","batch","normalization","over","a","batch","of","heterogeneous","features","as","described","in","the","`","``","Batch","Normalization",":","Accelerating","Deep","Network","Training","by","Reducing","Internal","Covariate","Shift","''","<","https",":","\/\/arxiv.org\/abs\/1502.03167",">","`","_","paper",".","Compared","to",":","class",":","`","BatchNorm","`",",",":","class",":","`","HeteroBatchNorm","`","applies","normalization","individually","for","each","node","or","edge","type",".","Args",":","in_channels","(","int",")",":","Size","of","each","input","sample",".","num_types","(","int",")",":","The","number","of","types",".","eps","(","float",",","optional",")",":","A","value","added","to","the","denominator","for","numerical","stability",".","(","default",":",":","obj",":","`","1e-5","`",")","momentum","(","float",",","optional",")",":","The","value","used","for","the","running","mean","and","running","variance","computation",".","(","default",":",":","obj",":","`","0.1","`",")","affine","(","bool",",","optional",")",":","If","set","to",":","obj",":","`","True","`",",","this","module","has","learnable","affine","parameters",":","math",":","`","\\gamma","`","and",":","math",":","`","\\beta","`",".","(","default",":",":","obj",":","`","True","`",")","track_running_stats","(","bool",",","optional",")",":","If","set","to",":","obj",":","`","True","`",",","this","module","tracks","the","running","mean","and","variance",",","and","when","set","to",":","obj",":","`","False","`",",","this","module","does","not","track","such","statistics","and","always","uses","batch","statistics","in","both","training","and","eval","modes",".","(","default",":",":","obj",":","`","True","`",")"]
docstr: Forward pass.
  
  Args:
      x (torch.Tensor): The input features.
      type_vec (torch.Tensor): A vector that maps each entry to a type.
docstr_tok: ["Forward","pass",".","Args",":","x","(","torch.Tensor",")",":","The","input","features",".","type_vec","(","torch.Tensor",")",":","A","vector","that","maps","each","entry","to","a","type","."]
code: def forward(self, x: Tensor, type_vec: Tensor) -> Tensor:
      r"""Forward pass.
  
      Args:
          x (torch.Tensor): The input features.
          type_vec (torch.Tensor): A vector that maps each entry to a type.
      """
      if not self.training and self.track_running_stats:
          mean, var = self.running_mean, self.running_var
      else:
          with torch.no_grad():
              mean, var = self.mean_var(
                  x, type_vec, dim_size=self.num_types
              )
  
      if self.training and self.track_running_stats:
          if self.momentum is None:
              self.num_batches_tracked.add_(1)
              exp_avg_factor = 1.0 / float(self.num_batches_tracked)
          else:
              exp_avg_factor = self.momentum
  
          with torch.no_grad():  # Update running mean and variance:
              type_index = torch.unique(type_vec)
  
              self.running_mean[type_index] = (
                  1.0 - exp_avg_factor
              ) * self.running_mean[type_index] + exp_avg_factor * mean[
                  type_index
              ]
              self.running_var[type_index] = (
                  1.0 - exp_avg_factor
              ) * self.running_var[type_index] + exp_avg_factor * var[
                  type_index
              ]
  
      out = (x - mean[type_vec]) / var.clamp(self.eps).sqrt()[type_vec]
  
      if self.affine:
          out = out * self.weight[type_vec] + self.bias[type_vec]
  
      return out
  
file_path: pytorch_geometric/torch_geometric/nn/norm/batch_norm.py
imports: from typing import Optional
  import torch
  from torch import Tensor
  from torch.nn import Parameter
  from torch_geometric.nn.aggr.fused import FusedAggregation
label_desc: Use image node_id 4 for calling the HeteroBatchNorm obj's underlying member method code with example usage: obj.forward(x, type_vec) and returns: out