mlx_graphs.nn.HeteroBatchNormalization

mlx_graphs.nn.HeteroBatchNormalization#

class mlx_graphs.nn.HeteroBatchNormalization(in_channels: int, num_types: int, eps: float = 1e-05, momentum: float | None = 0.1, affine: bool = True, track_running_stats: bool = True)[source]#

Bases: Module

Applies batch normalization over a batch of heterogeneous features, as described in the “Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift” paper.

Compared to a standard batch norm, HeteroBatchNorm 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)

  • momentum (float, optional) – The value used for the running mean and running variance computation. If set to None, a cumulative moving average is used instead. (default: 0.1)

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

  • track_running_stats (bool, optional) – If set to True, this module tracks the running mean and variance, and when set to False, this module does not track such statistics and always uses batch statistics in both training and eval modes. (default: True)

__call__(x: mlx.core.array, type_vec: mlx.core.array) → mlx.core.array[source]#

Forward pass.

Parameters:

  • x (mx.array) – The input features of shape [num_items, in_channels].

  • type_vec (mx.array) – An integer vector of shape [num_items] that maps each entry to a type in [0, num_types).

Return type:

array

Methods

reset_parameters()

Resets all learnable parameters of the module.

reset_running_stats()

Resets all running statistics of the module.
reset_parameters()[source]#

Resets all learnable parameters of the module.

reset_running_stats()[source]#

Resets all running statistics of the module.