TorchLean API

NN.Spec.Module.Normalization

Normalization module wrappers #

This file wraps selected normalization specs as Spec.Modules for composition/export.

For Spec.Module.layerNorm, the final axis is the feature axis and every preceding axis is treated as a batch axis. The implementation flattens those leading axes only while applying the matrix-level normalization semantics.

PyTorch mental picture: nn.LayerNorm(embedDim) applied at each timestep, with weight=gamma and bias=beta.

def Spec.Module.layerNorm {α : Type} [TorchLean.Storage α] [Context α] (leading : Shape) (width : ℕ) (gamma beta : TorchLean.Tensor α [width]) (hLeading : 0 < leading.size) (hWidth : 0 < width) (epsilon : α := TorchLean.normalizationEpsilon) :
Module α (leading.appendDim width) (leading.appendDim width)

LayerNorm over the final axis, wrapped as a Spec.Module.

The epsilon is passed unchanged to Spec.layerNorm; there is no representability check or fallback when it rounds to zero. Tiny formats may require an explicitly chosen positive finite epsilon. The Python expression records only the width, omitting the supplied epsilon and affine parameters.

Instances For