Skip to content

unconditioned_norm

unconditioned_norm

Wrapper that adapts plain norms to the conditioned normalization interface.

UnconditionedNorm

UnconditionedNorm(norm)

Bases: Module

Wraps a plain norm (LayerNorm, RMSNorm) to return (normed, gate).

Ignores the conditioning argument and returns gate = ones(1). This allows transformer blocks to call all normalizations uniformly.

Source code in src/versatil/models/layers/normalization/unconditioned_norm.py
def __init__(self, norm: nn.Module):
    super().__init__()
    self.norm = norm

forward

forward(x, condition=None)

Normalize the input and return it with a unit gate, ignoring the condition.

Source code in src/versatil/models/layers/normalization/unconditioned_norm.py
def forward(
    self,
    x: torch.Tensor,
    condition: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Normalize the input and return it with a unit gate, ignoring the condition."""
    normed = self.norm(x)
    gate = torch.ones(1, dtype=x.dtype, device=x.device)
    return normed, gate