BackwardScale

Backward (reverse-mode) scale propagation.

This optional module mirrors NN.Proofs.RuntimeApprox.Graph.BackwardApprox, but for scale bounds (nonnegative bounds on linfNorm) rather than eps error bounds.

Use it alongside the backward approximation graph when you want to derive abs+rel tolerances for gradients/cotangents from both an eps error bound and a propagated magnitude bound.

def Proofs.RuntimeApprox.AddScaleSound {α : Type} (toSpec : α → Spec.SpecScalar) [Add α] (addBound : {Δ : List Spec.Shape} → BList Δ → BList Δ → TorchLean.TensorPack α Δ → TorchLean.TensorPack α Δ → BList Δ) : Prop

Soundness condition for accumulating scale bounds under addition in a context.

structure Proofs.RuntimeApprox.RevNodeScale {α : Type} (toSpec : α → Spec.SpecScalar) (Γ : List Spec.Shape) (τ : Spec.Shape) extends Proofs.RuntimeApprox.RevNode toSpec Γ τ : Type

A reverse node augmented with forward+VJP scale bounds.

• forwardSpec : TorchLean.TensorPack Spec.SpecScalar Γ → Spec.SpecTensor τ
• forwardRuntime : TorchLean.TensorPack α Γ → TorchLean.Tensor α τ
• bound : EList Γ → TorchLean.TensorPack α Γ → Spec.SpecScalar
• sound (xS : TorchLean.TensorPack Spec.SpecScalar Γ) (xR : TorchLean.TensorPack α Γ) (eps : EList Γ) : approxCtx toSpec xS xR eps → approxTensor toSpec (self.forwardSpec xS) (self.forwardRuntime xR) (self.bound eps xR)
• vjpSpec : TorchLean.TensorPack Spec.SpecScalar Γ → Spec.SpecTensor τ → TorchLean.TensorPack Spec.SpecScalar Γ
• vjpRuntime : TorchLean.TensorPack α Γ → TorchLean.Tensor α τ → TorchLean.TensorPack α Γ
• vjpBound : EList Γ → TorchLean.TensorPack α Γ → Spec.SpecScalar → TorchLean.Tensor α τ → EList Γ
• vjpSound (ctxS : TorchLean.TensorPack Spec.SpecScalar Γ) (ctxR : TorchLean.TensorPack α Γ) (epsCtx : EList Γ) (δS : Spec.SpecTensor τ) (δR : TorchLean.Tensor α τ) (epsδ : Spec.SpecScalar) : approxCtx toSpec ctxS ctxR epsCtx → approxTensor toSpec δS δR epsδ → approxCtx toSpec (self.vjpSpec ctxS δS) (self.vjpRuntime ctxR δR) (self.vjpBound epsCtx ctxR epsδ δR)
• fwdScaleBound : BList Γ → TorchLean.TensorPack α Γ → NNReal
• fwdScaleSound (ctxS : TorchLean.TensorPack Spec.SpecScalar Γ) (ctxR : TorchLean.TensorPack α Γ) (epsCtx : EList Γ) (bCtx : BList Γ) : approxCtx toSpec ctxS ctxR epsCtx → scaleCtx toSpec ctxS ctxR bCtx → scaleTensor toSpec (self.forwardSpec ctxS) (self.forwardRuntime ctxR) (self.fwdScaleBound bCtx ctxR)
• vjpScaleBound : BList Γ → TorchLean.TensorPack α Γ → NNReal → TorchLean.Tensor α τ → BList Γ
• vjpScaleSound (ctxS : TorchLean.TensorPack Spec.SpecScalar Γ) (ctxR : TorchLean.TensorPack α Γ) (epsCtx : EList Γ) (bCtx : BList Γ) (δS : Spec.SpecTensor τ) (δR : TorchLean.Tensor α τ) (bδ : NNReal) : approxCtx toSpec ctxS ctxR epsCtx → scaleCtx toSpec ctxS ctxR bCtx → scaleTensor toSpec δS δR bδ → scaleCtx toSpec (self.vjpSpec ctxS δS) (self.vjpRuntime ctxR δR) (self.vjpScaleBound bCtx ctxR bδ δR)

inductive Proofs.RuntimeApprox.RevGraphScale {α : Type} (toSpec : α → Spec.SpecScalar) (Γ : List Spec.Shape) : List Spec.Shape → Type

Reverse-mode graph with scale-aware nodes.

• nil {α : Type} {toSpec : α → Spec.SpecScalar} {Γ : List Spec.Shape} : RevGraphScale toSpec Γ []
• snoc {α : Type} {toSpec : α → Spec.SpecScalar} {Γ ss : List Spec.Shape} {τ : Spec.Shape} : RevGraphScale toSpec Γ ss → RevNodeScale toSpec (Γ ++ ss) τ → RevGraphScale toSpec Γ (ss ++ [τ])

def Proofs.RuntimeApprox.RevGraphScale.toRevGraph {α : Type} {toSpec : α → Spec.SpecScalar} {Γ ss : List Spec.Shape} : RevGraphScale toSpec Γ ss → RevGraph toSpec Γ ss

Forget the scale annotations on nodes, producing an ordinary RevGraph.

def Proofs.RuntimeApprox.RevGraphScale.toFwdGraphScale {α : Type} {toSpec : α → Spec.SpecScalar} {Γ ss : List Spec.Shape} : RevGraphScale toSpec Γ ss → FwdGraphScale toSpec Γ ss

Convert a RevGraphScale into a FwdGraphScale by dropping the reverse-mode payload.

theorem Proofs.RuntimeApprox.RevGraphScale.toFwdGraph_toFwdGraphScale_eq {α : Type} {toSpec : α → Spec.SpecScalar} {Γ ss : List Spec.Shape} (g : RevGraphScale toSpec Γ ss) : g.toFwdGraphScale.toFwdGraph = g.toRevGraph.toFwdGraph

Forgetting the scale annotations commutes with forgetting the reverse-mode structure.

Both directions of the square end at the same plain forward graph, which is why the scale-annotated development can reuse the unannotated evaluation lemmas instead of repeating them.

def Proofs.RuntimeApprox.RevGraphScale.evalSpec {α : Type} {toSpec : α → Spec.SpecScalar} {Γ ss : List Spec.Shape} (g : RevGraphScale toSpec Γ ss) (x : TorchLean.TensorPack Spec.SpecScalar Γ) : TorchLean.TensorPack Spec.SpecScalar (Γ ++ ss)

Evaluate the forward pass on spec values, returning the extended context Γ ++ ss.

def Proofs.RuntimeApprox.RevGraphScale.evalRuntime {α : Type} {toSpec : α → Spec.SpecScalar} {Γ ss : List Spec.Shape} (g : RevGraphScale toSpec Γ ss) (x : TorchLean.TensorPack α Γ) : TorchLean.TensorPack α (Γ ++ ss)

Evaluate the forward pass on runtime values, returning the extended context Γ ++ ss.

@[simp]

theorem Proofs.RuntimeApprox.RevGraphScale.evalSpec_eq_rev {α : Type} {toSpec : α → Spec.SpecScalar} {Γ ss : List Spec.Shape} (g : RevGraphScale toSpec Γ ss) (x : TorchLean.TensorPack Spec.SpecScalar Γ) : g.evalSpec x = g.toRevGraph.evalSpec x

Spec evaluation of a scale-annotated reverse graph agrees with the unannotated one.

@[simp]

theorem Proofs.RuntimeApprox.RevGraphScale.evalRuntime_eq_rev {α : Type} {toSpec : α → Spec.SpecScalar} {Γ ss : List Spec.Shape} (g : RevGraphScale toSpec Γ ss) (x : TorchLean.TensorPack α Γ) : g.evalRuntime x = g.toRevGraph.evalRuntime x

And so does runtime evaluation. The annotations carry bounds only; they never change what is computed, which is exactly what these two lemmas record.

def Proofs.RuntimeApprox.RevGraphScale.evalBounds {α : Type} {toSpec : α → Spec.SpecScalar} {Γ ss : List Spec.Shape} (g : RevGraphScale toSpec Γ ss) (epsIn : EList Γ) (xR : TorchLean.TensorPack α Γ) : EList (Γ ++ ss)

Forward-pass error bounds for all intermediate nodes, computed from input bounds epsIn.

def Proofs.RuntimeApprox.RevGraphScale.evalScales {α : Type} {toSpec : α → Spec.SpecScalar} {Γ ss : List Spec.Shape} (g : RevGraphScale toSpec Γ ss) (bIn : BList Γ) (xR : TorchLean.TensorPack α Γ) : BList (Γ ++ ss)

Forward-pass scale bounds for all intermediate nodes, computed from input bounds bIn.

theorem Proofs.RuntimeApprox.RevGraphScale.eval_scale {α : Type} {toSpec : α → Spec.SpecScalar} {Γ ss : List Spec.Shape} (g : RevGraphScale toSpec Γ ss) (xS : TorchLean.TensorPack Spec.SpecScalar Γ) (xR : TorchLean.TensorPack α Γ) (epsIn : EList Γ) (bIn : BList Γ) : approxCtx toSpec xS xR epsIn → scaleCtx toSpec xS xR bIn → scaleCtx toSpec (g.evalSpec xS) (g.evalRuntime xR) (g.evalScales bIn xR)

Forward evaluation respects the scale bounds: if the inputs are within their bounds, every intermediate node is within the bound evalScales computes for it.

def Proofs.RuntimeApprox.RevGraphScale.backpropScales {α : Type} {toSpec : α → Spec.SpecScalar} {Γ ss : List Spec.Shape} (g : RevGraphScale toSpec Γ ss) [Add α] (bIn : BList Γ) (xR : TorchLean.TensorPack α Γ) (bSeed : BList (Γ ++ ss)) (seedR : TorchLean.TensorPack α (Γ ++ ss)) (addBound : {Δ : List Spec.Shape} → BList Δ → BList Δ → TorchLean.TensorPack α Δ → TorchLean.TensorPack α Δ → BList Δ) : BList Γ

Backpropagate scale bounds through a RevGraphScale, analogous to RevGraph.backpropRuntime.

theorem Proofs.RuntimeApprox.RevGraphScale.backprop_scale {α : Type} {toSpec : α → Spec.SpecScalar} {Γ ss : List Spec.Shape} (g : RevGraphScale toSpec Γ ss) [Add α] (addBound : {Δ : List Spec.Shape} → BList Δ → BList Δ → TorchLean.TensorPack α Δ → TorchLean.TensorPack α Δ → BList Δ) (addSound : AddScaleSound toSpec fun {Δ : List Spec.Shape} => addBound) (xS : TorchLean.TensorPack Spec.SpecScalar Γ) (xR : TorchLean.TensorPack α Γ) (epsIn : EList Γ) (bIn : BList Γ) (seedS : TorchLean.TensorPack Spec.SpecScalar (Γ ++ ss)) (seedR : TorchLean.TensorPack α (Γ ++ ss)) (bSeed : BList (Γ ++ ss)) : approxCtx toSpec xS xR epsIn → scaleCtx toSpec xS xR bIn → scaleCtx toSpec seedS seedR bSeed → scaleCtx toSpec (g.toRevGraph.backpropSpec xS seedS) (g.toRevGraph.backpropRuntime xR seedR) (g.backpropScales bIn xR bSeed seedR fun {Δ : List Spec.Shape} => addBound)

The same statement for the backward pass: propagated cotangents stay within the backpropagated scale bounds.

The addBound combiner is a parameter rather than a fixed choice, because how two accumulated gradients' bounds combine depends on the carrier: adding bounds is always sound, but a carrier with a sharper triangle inequality can do better. addSound is what pins down the requirement.