DAG Normalization Primitives

RMS normalization and regularized L2 normalization with a nonempty final axis. The L2 epsilon is an unchecked scalar input; callers choose a positive value.

def NN.GraphSpec.DAG.PrimOp.Internal.mapFinalAxis {α : Type} [TorchLean.Storage α] (leading : Spec.Shape) {width : ℕ} (f : TorchLean.Tensor α [width] → TorchLean.Tensor α [width]) (input : TorchLean.Tensor α (leading.appendDim width)) : TorchLean.Tensor α (leading.appendDim width)

Adapt the shared leading-axis traversal to the final-axis shape convention.

def NN.GraphSpec.DAG.PrimOp.Internal.rmsNormVectorSpec {α : Type} [TorchLean.Storage α] [Context α] {width : ℕ} (hWidth : 0 < width) (input gamma : TorchLean.Tensor α [width]) : TorchLean.Tensor α [width]

RMS normalization of a single vector.

The spec is written for a batch, so a lone vector is normalized by wrapping it in a one-element batch and reading the result back out. hWidth is what rules out dividing by an empty mean.

def NN.GraphSpec.DAG.PrimOp.rmsNormSpec {α : Type} [TorchLean.Storage α] [Context α] (leading : Spec.Shape) {width : ℕ} (hWidth : 0 < width) (gamma : TorchLean.Tensor α [width]) : TorchLean.Tensor α (leading.appendDim width) → TorchLean.Tensor α (leading.appendDim width)

Apply RMS normalization with a shared final-axis scale over an arbitrary leading shape.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.rmsNormSpec_scalar {α : Type} [TorchLean.Storage α] [Context α] {width : ℕ} (hWidth : 0 < width) (gamma input : TorchLean.Tensor α [width]) : rmsNormSpec Spec.Shape.scalar hWidth gamma input = Internal.rmsNormVectorSpec hWidth input gamma

With no leading axes, RMS normalization uses the single-vector spec.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.rmsNormSpec_dim {α : Type} [TorchLean.Storage α] [Context α] {count width : ℕ} {rest : Spec.Shape} (hWidth : 0 < width) (gamma : TorchLean.Tensor α [width]) (values : Fin count → TorchLean.Tensor α (rest.appendDim width)) : rmsNormSpec (Spec.Shape.dim count rest) hWidth gamma (TorchLean.Tensor.dim values) = TorchLean.Tensor.dim fun (index : Fin count) => rmsNormSpec rest hWidth gamma (values index)

Leading axes distribute over RMS normalization without changing the row computation.

def NN.GraphSpec.DAG.PrimOp.rmsNorm (leading : Spec.Shape) (width : ℕ) (hWidth : 0 < width) : PrimOp [leading.appendDim width, [width]] (leading.appendDim width)

Root-mean-square normalization along the final axis of an arbitrary tensor.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.rmsNorm_specFwd {leading : Spec.Shape} {width : ℕ} (hWidth : 0 < width) {α : Type} [TorchLean.Storage α] [Context α] (input : TorchLean.Tensor α (leading.appendDim width)) (gamma : TorchLean.Tensor α [width]) : (rmsNorm leading width hWidth).specFwd (TorchLean.TensorPack.cons input (TorchLean.TensorPack.cons gamma TorchLean.TensorPack.nil)) = rmsNormSpec leading hWidth gamma input

Pure evaluation of final-axis RMS normalization with a shared scale.

def NN.GraphSpec.DAG.PrimOp.normalizeL2RegularizedLastAxisSpec {α : Type} [TorchLean.Storage α] [Context α] (leading : Spec.Shape) {width : ℕ} (epsilon : α) : TorchLean.Tensor α (leading.appendDim width) → TorchLean.Tensor α (leading.appendDim width)

Apply regularized L2 normalization over the final axis of an arbitrary tensor.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.normalizeL2RegularizedLastAxisSpec_scalar {α : Type} [TorchLean.Storage α] [Context α] {width : ℕ} (epsilon : α) (values : TorchLean.Tensor α [width]) : normalizeL2RegularizedLastAxisSpec Spec.Shape.scalar epsilon values = Spec.normalizeL2RegularizedSpec values epsilon

With no batch axes left, the graph node is the regularized L2 normalization spec.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.normalizeL2RegularizedLastAxisSpec_dim {α : Type} [TorchLean.Storage α] [Context α] {count width : ℕ} {rest : Spec.Shape} (epsilon : α) (values : Fin count → TorchLean.Tensor α (rest.appendDim width)) : normalizeL2RegularizedLastAxisSpec (Spec.Shape.dim count rest) epsilon (TorchLean.Tensor.dim values) = TorchLean.Tensor.dim fun (index : Fin count) => normalizeL2RegularizedLastAxisSpec rest epsilon (values index)

Batch axes distribute over L2 normalization.

def NN.GraphSpec.DAG.PrimOp.normalizeL2Regularized (leading : Spec.Shape) (width : ℕ) (hWidth : 0 < width) : PrimOp [leading.appendDim width, Spec.Shape.scalar] (leading.appendDim width)

Normalize the final axis by sqrt(sum (x * x) + epsilon).

This is not torch.nn.functional.normalize, which divides by max(‖x‖, epsilon). The two agree for nonzero real vectors when epsilon = 0; positive epsilon changes the denominator in general.

The stabilizer is a scalar graph input rather than a declaration parameter. This keeps the exact numerical convention visible in captured graphs and permits architectures to select their own positive epsilon without adding another primitive.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.normalizeL2Regularized_specFwd {leading : Spec.Shape} {width : ℕ} (hWidth : 0 < width) {α : Type} [TorchLean.Storage α] [Context α] (input : TorchLean.Tensor α (leading.appendDim width)) (epsilon : α) : (normalizeL2Regularized leading width hWidth).specFwd (TorchLean.TensorPack.cons input (TorchLean.TensorPack.cons (TorchLean.Tensor.scalar epsilon) TorchLean.TensorPack.nil)) = normalizeL2RegularizedLastAxisSpec leading epsilon input

Pure evaluation of final-axis regularized L2 normalization.