factory

factory

Normalization layers factory method.

create_normalization_layer

create_normalization_layer(normalization_type, dimension, epsilon=1e-06, conditioning_dimension=None)

Create a normalization layer, optionally wrapped with adaptive conditioning.

When conditioning_dimension is provided, returns an AdaNorm that wraps the base normalization with a learned modulation. Otherwise returns a plain norm.

Parameters:

Name	Type	Description	Default
normalization_type	str	Base normalization type (use NormalizationType enum values).	required
dimension	int	Feature dimension.	required
epsilon	float	Small constant for numerical stability.	1e-06
conditioning_dimension	int | None	Conditioning dimension. When set, wraps the base norm in AdaNorm for adaptive modulation.	None

Returns:

Type	Description
Module	Plain normalization layer or AdaNorm.

Raises:

Type	Description
ValueError	If normalization_type is not supported.

Source code in src/versatil/models/layers/normalization/factory.py

def create_normalization_layer(
    normalization_type: str,
    dimension: int,
    epsilon: float = 1e-6,
    conditioning_dimension: int | None = None,
) -> nn.Module:
    """Create a normalization layer, optionally wrapped with adaptive conditioning.

    When ``conditioning_dimension`` is provided, returns an AdaNorm that wraps the base
    normalization with a learned modulation. Otherwise returns a plain norm.

    Args:
        normalization_type: Base normalization type (use NormalizationType enum values).
        dimension: Feature dimension.
        epsilon: Small constant for numerical stability.
        conditioning_dimension: Conditioning dimension. When set, wraps the base norm
            in AdaNorm for adaptive modulation.

    Returns:
        Plain normalization layer or AdaNorm.

    Raises:
        ValueError: If normalization_type is not supported.
    """
    if conditioning_dimension is not None:
        base_norm = _create_base_norm(
            normalization_type=normalization_type,
            dimension=dimension,
            epsilon=epsilon,
            elementwise_affine=False,
        )
        return AdaNorm(
            base_norm=base_norm,
            conditioning_dimension=conditioning_dimension,
            feature_dim=dimension,
        )
    return _create_base_norm(
        normalization_type=normalization_type,
        dimension=dimension,
        epsilon=epsilon,
    )

create_block_normalization

create_block_normalization(normalization_type, dimension, epsilon=1e-06, conditioning_dimension=None, use_gating=False, init_strategy='zero')

Create normalization for transformer blocks: (x, condition) -> (normed, gate).

When conditioning_dimension is provided, returns an AdaNorm with learned modulation (and optional gating for AdaLN-Zero). Otherwise returns an UnconditionedNorm that wraps a plain normalization layer.

Parameters:

Name	Type	Description	Default
normalization_type	str	Base normalization type (use NormalizationType enum values).	required
dimension	int	Feature dimension.	required
epsilon	float	Small constant for numerical stability.	1e-06
conditioning_dimension	int | None	Conditioning dimension. When set, creates AdaNorm.	None
use_gating	bool	Whether to produce a learned gate (AdaLN-Zero). Only applies when conditioning_dimension is set.	False
init_strategy	Literal['zero', 'xavier']	Initialization strategy for modulation weights.	'zero'

Returns:

Type	Description
BlockNormalization	AdaNorm when conditioned, UnconditionedNorm when not.

Raises:

Type	Description
ValueError	If normalization_type is not supported.

Source code in src/versatil/models/layers/normalization/factory.py

def create_block_normalization(
    normalization_type: str,
    dimension: int,
    epsilon: float = 1e-6,
    conditioning_dimension: int | None = None,
    use_gating: bool = False,
    init_strategy: Literal["zero", "xavier"] = "zero",
) -> BlockNormalization:
    """Create normalization for transformer blocks: ``(x, condition) -> (normed, gate)``.

    When ``conditioning_dimension`` is provided, returns an AdaNorm with learned
    modulation (and optional gating for AdaLN-Zero). Otherwise returns
    an UnconditionedNorm that wraps a plain normalization layer.

    Args:
        normalization_type: Base normalization type (use NormalizationType enum values).
        dimension: Feature dimension.
        epsilon: Small constant for numerical stability.
        conditioning_dimension: Conditioning dimension. When set, creates AdaNorm.
        use_gating: Whether to produce a learned gate (AdaLN-Zero).
            Only applies when conditioning_dimension is set.
        init_strategy: Initialization strategy for modulation weights.

    Returns:
        AdaNorm when conditioned, UnconditionedNorm when not.

    Raises:
        ValueError: If normalization_type is not supported.
    """
    if conditioning_dimension is not None:
        base_norm = _create_base_norm(
            normalization_type=normalization_type,
            dimension=dimension,
            epsilon=epsilon,
            elementwise_affine=False,
        )
        return AdaNorm(
            base_norm=base_norm,
            conditioning_dimension=conditioning_dimension,
            feature_dim=dimension,
            use_gate=use_gating,
            init_strategy=init_strategy,
        )
    base_norm = _create_base_norm(
        normalization_type=normalization_type,
        dimension=dimension,
        epsilon=epsilon,
    )
    return UnconditionedNorm(base_norm)