LayerNorm IR Evaluation

The IR layernorm axis node interprets axis as the start of the normalized suffix. Evaluation reshapes the tensor to (seqLen, embedDim), runs the spec 2D LayerNorm with gamma=1 and beta=0, then reshapes back.

theorem NN.Verification.Builtin.Proved.Correctness.IRStep.layerNormWithoutAffine_eq_spec {α : Type} [TorchLean.Storage α] [Context α] (seqLen embedDim : ℕ) (x : TorchLean.Tensor α [seqLen, embedDim]) (hSeq : seqLen > 0) (hEmb : embedDim > 0) : IR.Graph.layerNormWithoutAffine seqLen embedDim x = Except.ok (Spec.layerNorm x (TorchLean.Tensor.full (Spec.Shape.dim embedDim Spec.Shape.scalar) 1) (TorchLean.Tensor.full (Spec.Shape.dim embedDim Spec.Shape.scalar) 0) hSeq hEmb)

Affine-free IR LayerNorm is exactly spec LayerNorm with unit scale and zero bias.

theorem NN.Verification.Builtin.Proved.Correctness.IRStep.evalAt_layernorm_eq {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (axis seqLen embedDim : ℕ) (x : TorchLean.Tensor α s) (hParams : IR.OpContracts.layerNormMatrixDims axis s = Except.ok (seqLen, embedDim)) (hNumel : s.size = (Spec.Shape.dim seqLen (Spec.Shape.dim embedDim Spec.Shape.scalar)).size) (y2d : TorchLean.Tensor α [seqLen, embedDim]) (hLayerNorm : IR.Graph.layerNormWithoutAffine seqLen embedDim (x.reshapeSpec hNumel) = Except.ok y2d) : (unaryGraphOut (IR.OpKind.layernorm axis) s s).evalAt { } { shape := s, tensor := x } #[{ shape := s, tensor := x }] 1 = Except.ok { shape := s, tensor := y2d.reshapeSpec ⋯ }

Local IR semantics for layernorm axis.

The hypotheses are the same contracts checked by the evaluator: axis must produce a valid (seqLen, embedDim) view, the reshape must preserve element count, and the affine-free LayerNorm step must succeed.