norms

Normalization factory wrapper with spatial dimensionality support.

PyG's normalization_resolver covers graph-aware norms and nn.BatchNorm1d but does not register the nn.BatchNorm2d / nn.BatchNorm3d variants needed by convolutional stacks. create_norm adds a dim parameter that selects the matching nn.*Nd class for the common families when dim > 1, and defers to PyG's resolver otherwise.

Functions:

create_norm –

Resolve a normalization layer with a spatial dimensionality hint.

create_norm ¶

create_norm(
    norm: Union[str, Callable, None],
    channels: int,
    *,
    dim: int = 1,
    conditions: Optional[Sequence[str]] = None,
    **norm_kwargs: Any,
) -> Optional[Module]

Resolve a normalization layer with a spatial dimensionality hint.

For dim == 1, defers to PyG's normalization_resolver (graph-aware norms plus BatchNorm1d). For dim \(\in \{2, 3\}\), maps common norm names to the matching nn.*Nd variant. Pass a class or an existing instance to bypass string resolution. When conditions is given, the resolved norm is wrapped in a per-condition PDNorm for multi-dataset (prompt-driven) training.

Parameters:

norm (Union[str, Callable, None]) –

Norm name ("batch_norm", "instance_norm", "group_norm", "layer_norm"), a class, an instance, or None.
channels (int) –

Number of feature channels.
dim (int, default: 1 ) –

Spatial dimensionality of the feature map. \(1\) for packed / graph tensors \((N, C)\), \(2\) for \((B, C, H, W)\), \(3\) for \((B, C, H, W, D)\).
conditions (Optional[Sequence[str]], default: None ) –

Ordered dataset condition names. When set (and norm is not None), returns a PDNorm holding one inner norm per condition.
**norm_kwargs (Any, default: {} ) –

Forwarded to the norm constructor.

Returns:

Optional[Module] –

The instantiated norm module, or None if norm is None.