torch_geometric.nn.norm.HeteroLayerNorm

class HeteroLayerNorm(in_channels: int, num_types: int, eps: float = 1e-05, affine: bool = True, mode: str = 'node')[source]

Bases: Module

Applies layer normalization over each individual example in a batch of heterogeneous features as described in the “Layer Normalization” paper. Compared to LayerNorm, HeteroLayerNorm applies normalization individually for each node or edge type.

Parameters:
  • 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: 1e-5)

  • affine (bool, optional) – If set to True, this module has learnable affine parameters \(\gamma\) and \(\beta\). (default: True)

  • mode (str, optinal) – The normalization mode to use for layer normalization ("node"). If “node” is used, each node will be considered as an element to be normalized. (default: "node")

reset_parameters()[source]

Resets all learnable parameters of the module.

forward(x: Tensor, type_vec: Optional[Tensor] = None, type_ptr: Optional[Union[Tensor, List[int]]] = None) Tensor[source]

Forward pass.

Note

Either type_vec or type_ptr needs to be specified. In general, relying on type_ptr is more efficient in case the input tensor is sorted by types.

Parameters:
  • x (torch.Tensor) – The input features.

  • type_vec (torch.Tensor, optional) – A vector that maps each entry to a type. (default: None)

  • type_ptr (torch.Tensor or List[int]) – A vector denoting the boundaries of types. (default: None)