Dense Backward Pass: Appending One Lowered Node

This file proves the inductive step shared by the Graph and GraphData backward-link theorems. Given a prefix tape t on which the dense reverse loop is already known to agree with some reverse program bp, appending one lowered node yields a tape on which the dense reverse loop computes: one vjp step for the new node, added into the prefix seed, followed by bp on the prefix. The last gradient slot is never modified afterwards.

The proof is organised in three groups of lemmas.

• Push commutation: addGradAll, backwardDenseFromStep and backwardDenseFromLoop on the extended tape and an accumulator with one extra slot act on the prefix tape and accumulator, leaving the extra slot untouched, as long as all ids stay in the prefix range.
• The last-node step: processing the freshly appended node adds its vjp contributions into the prefix, which is TensorPack.add after shape erasure.
• backwardDenseFrom_addNode_lowerNode, which assembles the two.

Everything is stated for an arbitrary prefix tape characterised by its size, its requiresGrad flags, its stored values and the BackwardPidsLt invariant, so it does not depend on which graph representation produced the prefix, nor on the debug name of the appended node.

theorem Proofs.Autograd.Algebra.Graph.getNode?_addNode_of_lt {α : Type} [TorchLean.Storage α] (t : Runtime.Autograd.Tape α) (nd : Runtime.Autograd.Node α) {id : ℕ} (hid : id < t.nodes.size) : (t.addNode nd).1.getNode? id = t.getNode? id

Appending a node does not change the nodes stored at existing ids.

theorem Proofs.Autograd.Algebra.Graph.addGradAll_addNode_push {α : Type} [TorchLean.Storage α] [Add α] (t : Runtime.Autograd.Tape α) (nd : Runtime.Autograd.Node α) (v : Spec.SomeTensor α) (acc : Array (Spec.SomeTensor α)) (hacc : acc.size = t.nodes.size) (pid : ℕ) (pg : Spec.SomeTensor α) (hpid : pid < t.nodes.size) : (t.addNode nd).1.addGradAll (acc.push v) pid pg = Except.map (fun (a : Array (Spec.SomeTensor α)) => a.push v) (t.addGradAll acc pid pg)

addGradAll commutes with pushing an unused last slot.

On the tape extended by nd and an accumulator extended by v, accumulating into an id of the prefix range acts exactly as on the prefix tape and accumulator, and v is carried along.

theorem Proofs.Autograd.Algebra.Graph.foldlM_addGradAll_addNode_push {α : Type} [TorchLean.Storage α] [Add α] (t : Runtime.Autograd.Tape α) (nd : Runtime.Autograd.Node α) (v : Spec.SomeTensor α) (cs : List (ℕ × Spec.SomeTensor α)) (acc : Array (Spec.SomeTensor α)) : acc.size = t.nodes.size → (∀ {pid : ℕ} {pg : Spec.SomeTensor α}, (pid, pg) ∈ cs → pid < t.nodes.size) → List.foldlM (fun (acc2 : Array (Spec.SomeTensor α)) (x : ℕ × Spec.SomeTensor α) => match x with | (pid, pg) => (t.addNode nd).1.addGradAll acc2 pid pg) (acc.push v) cs = Except.map (fun (a : Array (Spec.SomeTensor α)) => a.push v) (List.foldlM (fun (acc2 : Array (Spec.SomeTensor α)) (x : ℕ × Spec.SomeTensor α) => match x with | (pid, pg) => t.addGradAll acc2 pid pg) acc cs)

Folding addGradAll over a contribution list commutes with pushing an unused last slot, as long as every contribution targets the prefix range.

theorem Proofs.Autograd.Algebra.Graph.backwardDenseFromStep_addNode_push {α : Type} [TorchLean.Storage α] [Add α] (t : Runtime.Autograd.Tape α) (nd : Runtime.Autograd.Node α) (v : Spec.SomeTensor α) (hpids : BackwardPidsLt t) {id : ℕ} (hid : id < t.nodes.size) (acc : Array (Spec.SomeTensor α)) (hacc : acc.size = t.nodes.size) : (t.addNode nd).1.backwardDenseFromStep (acc.push v) id = Except.map (fun (a : Array (Spec.SomeTensor α)) => a.push v) (t.backwardDenseFromStep acc id)

backwardDenseFromStep commutes with pushing an unused last slot for ids in the prefix range.

The BackwardPidsLt invariant guarantees that all contributions emitted by the visited node also stay in the prefix range.

theorem Proofs.Autograd.Algebra.Graph.backwardDenseFromLoop_addNode_push {α : Type} [TorchLean.Storage α] [Add α] (t : Runtime.Autograd.Tape α) (nd : Runtime.Autograd.Node α) (v : Spec.SomeTensor α) (hpids : BackwardPidsLt t) (m : ℕ) : m ≤ t.nodes.size → ∀ (acc : Array (Spec.SomeTensor α)), acc.size = t.nodes.size → (t.addNode nd).1.backwardDenseFromLoop m (acc.push v) = Except.map (fun (a : Array (Spec.SomeTensor α)) => a.push v) (t.backwardDenseFromLoop m acc)

backwardDenseFromLoop over the prefix range commutes with pushing an unused last slot.

theorem Proofs.Autograd.Algebra.Graph.addNode_prefix_node_spec {α : Type} [TorchLean.Storage α] (t : Runtime.Autograd.Tape α) (nd : Runtime.Autograd.Node α) {ss : List Spec.Shape} (ctx seed : TorchLean.TensorPack α ss) (hsize : t.nodes.size = ss.length) (hreq : ∀ (i : ℕ) (hi : i < t.nodes.size), t.nodes[i].requiresGrad = true) (hvals : Array.map (fun (n : Runtime.Autograd.Node α) => n.value) t.nodes = ctx.toShapeErasedArray) (i : ℕ) (hi : i < ss.length) : ∃ (node : Runtime.Autograd.Node α), (t.addNode nd).1.getNode? i = some node ∧ node.requiresGrad = true ∧ node.value.shape = seed.toShapeErasedArray[i].shape

The prefix of the extended tape is a well-formed dense accumulator for seed: every prefix id stores a node that requires gradients and whose value has the shape of the seed entry.

This is the hypothesis of foldlM_addGradAll_toIndexedShapeErasedArray_eq_add with pref = #[].

theorem Proofs.Autograd.Algebra.Graph.backwardDenseFromStep_addNode_lowerNode_last {α : Type} [TorchLean.Storage α] [Add α] {Δ : Type} {ss : List Spec.Shape} {τ : Spec.Shape} (t : Runtime.Autograd.Tape α) (name : Option String) (node : NodeData α Δ ss τ) (ctx : TorchLean.TensorPack α ss) (d : Δ) (seedPrev : TorchLean.TensorPack α ss) (seedOut : TorchLean.Tensor α τ) (hsize : t.nodes.size = ss.length) (hreq : ∀ (i : ℕ) (hi : i < t.nodes.size), t.nodes[i].requiresGrad = true) (hvals : Array.map (fun (n : Runtime.Autograd.Node α) => n.value) t.nodes = ctx.toShapeErasedArray) : (t.addNode (lowerNode name node ctx d)).1.backwardDenseFromStep (seedPrev.toShapeErasedArray.push (Spec.SomeTensor.ofTensor seedOut)) t.nodes.size = Except.ok ((seedPrev.add (node.vjp ctx d seedOut)).toShapeErasedArray.push (Spec.SomeTensor.ofTensor seedOut))

The runtime step for a freshly appended lowered node adds the node's vjp contributions into the prefix gradients, which is TensorPack.add after shape erasure, and leaves the node's own slot untouched.

theorem Proofs.Autograd.Algebra.Graph.backwardDenseFrom_addNode_lowerNode {α : Type} [TorchLean.Storage α] [Add α] {Δ : Type} {ss : List Spec.Shape} {τ : Spec.Shape} (t : Runtime.Autograd.Tape α) (name : Option String) (node : NodeData α Δ ss τ) (ctx : TorchLean.TensorPack α ss) (d : Δ) (bp : TorchLean.TensorPack α ss → TorchLean.TensorPack α ss) (hsize : t.nodes.size = ss.length) (hreq : ∀ (i : ℕ) (hi : i < t.nodes.size), t.nodes[i].requiresGrad = true) (hvals : Array.map (fun (n : Runtime.Autograd.Node α) => n.value) t.nodes = ctx.toShapeErasedArray) (hpids : BackwardPidsLt t) (hbp : ∀ (s : TorchLean.TensorPack α ss), t.backwardDenseFrom s.toShapeErasedArray = Except.ok (bp s).toShapeErasedArray) (seedPrev : TorchLean.TensorPack α ss) (seedOut : TorchLean.Tensor α τ) : (t.addNode (lowerNode name node ctx d)).1.backwardDenseFrom (seedPrev.snoc seedOut).toShapeErasedArray = Except.ok ((bp (seedPrev.add (node.vjp ctx d seedOut))).snoc seedOut).toShapeErasedArray

Inductive step of the backward link. Let t be a tape whose dense reverse loop agrees with the reverse program bp on every seed (after shape erasure), and whose nodes have the sizes, flags, values and backward-pointing invariant of a lowered prefix over the context ss. Appending the lowered node node yields a tape whose dense reverse loop, seeded by snoc seedPrev seedOut, returns snoc (bp (seedPrev + node.vjp ctx d seedOut)) seedOut.

This is the shared content of backwardDenseFrom_lowerGraphToTape_eq_backpropAllCtx and its GraphData variant: the snoc cases of backpropAllCtx unfold to exactly this shape.