mlx_graphs.nn.HeteroLayerNormalization

class mlx_graphs.nn.HeteroLayerNormalization(in_channels: int, num_types: int, eps: float = 1e-05, affine: bool = True)

Bases: Module

Applies layer normalization over each individual example in a batch of heterogeneous features, as described in the “Layer Normalization” paper.

Compared to a standard layer norm, HeteroLayerNorm applies normalization individually for each node or edge type. Note that the normalization itself (per-node mean/variance across channels) does not depend on the type: only the learnable affine parameters do.

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)

__call__(x: mlx.core.array, type_vec: mlx.core.array) → mlx.core.array

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_parameters()

Resets all learnable parameters of the module.