The eager LayerNorm rule

The eager tape checks the shapes of its input, scale, and bias, then saves those tensors in one LayerNorm node. Its backward closure calls Spec.layerNormBackward and tags each cotangent with the corresponding parent id.

These theorems inspect that recorded node. The forward value is the actual Spec.layerNorm output, and applying its backward closure to a correctly shaped cotangent returns the three blocks of the adjoint Fréchet derivative. The input-read hypotheses describe successful dynamic shape checks; the derivative identity is proved from the real LayerNorm formula.

Parent ids may coincide. The closure still returns one contribution per argument; accumulation of contributions into a shared parent belongs to the tape traversal. This file concerns the exact real instantiation of the eager tape, whose tensor arithmetic has mathematical semantics.

theorem Proofs.Autograd.LayerNorm.eager_layerNorm_value {m n : ℕ} (hm : 0 < m) (hn : 0 < n) (t : Runtime.Autograd.Tape ℝ) (xId gammaId betaId : ℕ) (x : TorchLean.Tensor ℝ [m, n]) (gamma beta : TorchLean.Tensor ℝ [n]) (ε : ℝ) (hx : t.requireValue xId = Except.ok x) (hgamma : t.requireValue gammaId = Except.ok gamma) (hbeta : t.requireValue betaId = Except.ok beta) : ((Except.toOption (Runtime.Autograd.Tape.layerNorm hm hn t xId gammaId betaId ε)).bind fun (result : Runtime.Autograd.Tape ℝ × ℕ) => result.1.getValue? result.2) = some (Spec.SomeTensor.ofTensor (Spec.layerNorm x gamma beta hm hn ε))

Recording LayerNorm saves its specified forward value at the returned node id.

The successful lookups retain the original scale and bias tensors and the configured epsilon.

theorem Proofs.Autograd.LayerNorm.eager_layerNorm_backward_eq_adjoint_fderiv {m n : ℕ} (hm : 0 < m) (hn : 0 < n) {ε : ℝ} (hε : 0 < ε) (t : Runtime.Autograd.Tape ℝ) (xId gammaId betaId : ℕ) (x gradOutput : TorchLean.Tensor ℝ [m, n]) (gamma beta : TorchLean.Tensor ℝ [n]) (hx : t.requireValue xId = Except.ok x) (hgamma : t.requireValue gammaId = Except.ok gamma) (hbeta : t.requireValue betaId = Except.ok beta) : ((Except.toOption (Runtime.Autograd.Tape.layerNorm hm hn t xId gammaId betaId ε)).bind fun (result : Runtime.Autograd.Tape ℝ × ℕ) => Option.map (fun (node : Runtime.Autograd.Node ℝ) => node.backward (Spec.SomeTensor.ofTensor gradOutput)) (result.1.getNode? result.2)) = have gradient := (ContinuousLinearMap.adjoint (fderiv ℝ (specLayerNormVec hm hn ε) (packLN x gamma beta))) (tensorToVec gradOutput); some (Except.ok #[(xId, Spec.SomeTensor.ofTensor (specX gradient)), (gammaId, Spec.SomeTensor.ofTensor (specGamma gradient)), (betaId, Spec.SomeTensor.ofTensor (specBeta gradient))])

The stored eager backward closure returns the adjoint derivative in input, scale, bias order.

Reading the just-recorded node and invoking its closure both succeed. The resulting array keeps the runtime parent ids, so the equality also checks the correspondence between argument slots and cotangent destinations.