TorchLean API

NN.Proofs.Autograd.Runtime.Link.BackwardDense

Executed Backward Pass versus Proved Backward Pass #

The eager trainer executes Tape.backwardDenseAll, which runs Tape.backwardDense: an Option-valued reverse sweep that runs a node's VJP only when the node has received a cotangent, then zero-fills the unreached slots. The link theorems in BackwardGraph and FDeriv are about Tape.backwardDenseFrom, which starts from a total gradient array and runs every VJP.

This file closes that gap at the tape level. The hypothesis is ZeroPreserving: every VJP on the tape sends the node's zero cotangent to zero contributions of the parents' shapes. Under it,

backwardDenseAll_eq_backwardDenseFrom : Tape.backwardDenseAll t outId seed = Tape.backwardDenseFrom t (oneHotGrads t outId seed)

holds for any tape, any valid output id and any seed of the output's shape, including the error cases. The proof is a step-by-step simulation: Tape.totalize maps the optional accumulator of backwardDense to the total accumulator of backwardDenseFrom, and each reverse step commutes with it. A skipped (unreached) node on the executed side corresponds to a VJP call on the zero cotangent on the proved side, which by ZeroPreserving adds only zeros.

BackwardDenseGraph proves that every tape produced by lowerGraphToTape is ZeroPreserving and derives the corollaries for the executed backward pass.

Zero cotangents and totalization #

The zero cotangent of a runtime node: an all-zero tensor of the node's value shape.

Instances For
    @[simp]

    A zero cotangent has the same shape as the node's value, by construction.

    The initial total gradient array of backwardDenseFrom: seed at outId, zeros elsewhere.

    Instances For

      backwardDenseAll is backwardDense followed by Tape.totalize.

      @[simp]

      Totalization produces exactly one entry per tape node.

      @[simp]

      So does the one-hot seed array, which is why both can be indexed by node id without bounds checks in the invariants below.

      theorem Proofs.Autograd.Algebra.Graph.getElem?_totalize {α : Type} [TorchLean.Storage α] [Zero α] (t : Runtime.Autograd.Tape (Spec.SomeTensor α)) (grads : Array (Option (Spec.SomeTensor α))) (id : ℕ) :
      (t.totalize grads)[id]? = Option.map (fun (node : Runtime.Autograd.Node (Spec.SomeTensor α)) => match grads[id]? with | some (some grad) => grad | x => zeroCotangent node) (t.getNode? id)

      Entry id of the totalized array, expressed through getNode?.

      theorem Proofs.Autograd.Algebra.Graph.getElem?_oneHotGrads {α : Type} [TorchLean.Storage α] [Zero α] (t : Runtime.Autograd.Tape (Spec.SomeTensor α)) (outId : ℕ) (seed : Spec.SomeTensor α) (id : ℕ) :
      (oneHotGrads t outId seed)[id]? = Option.map (fun (node : Runtime.Autograd.Node (Spec.SomeTensor α)) => if id = outId then seed else zeroCotangent node) (t.getNode? id)

      Entry id of the one-hot array, expressed through getNode?.

      theorem Proofs.Autograd.Algebra.Graph.getElem?_totalize_of_some {α : Type} [TorchLean.Storage α] [Zero α] {t : Runtime.Autograd.Tape (Spec.SomeTensor α)} {grads : Array (Option (Spec.SomeTensor α))} {id : ℕ} {node : Runtime.Autograd.Node (Spec.SomeTensor α)} {g : Spec.SomeTensor α} (hnode : t.getNode? id = some node) (hg : grads[id]? = some (some g)) :
      (t.totalize grads)[id]? = some g

      A node that the backward pass reached keeps its computed gradient.

      A node the backward pass never reached gets its zero cotangent.

      Taken together, these two lemmas say totalization loses no information: none in the sparse array means the node genuinely received no contribution, so filling in zero is not an approximation.

      theorem Proofs.Autograd.Algebra.Graph.totalize_set {α : Type} [TorchLean.Storage α] [Zero α] (t : Runtime.Autograd.Tape (Spec.SomeTensor α)) (grads : Array (Option (Spec.SomeTensor α))) (id : ℕ) (g : Spec.SomeTensor α) (hid : id < grads.size) (hidt : id < t.nodes.size) :
      t.totalize (grads.set id (some g) hid) = (t.totalize grads).set id g ⋯

      Totalization commutes with writing a present gradient.

      Invariants #

      Well-formedness of the optional accumulator used by backwardDense: one slot per tape node, and every present gradient has its node's value shape.

      Instances For

        Well-formedness of the total accumulator used by backwardDenseFrom: one slot per tape node, and every slot has its node's value shape.

        Instances For

          Zero preservation of a tape: every node's VJP sends the node's zero cotangent to an array of contributions each of which is the zero cotangent of an existing parent node.

          This is the exact condition under which skipping unreached nodes (backwardDense) and running every node (backwardDenseFrom) produce the same gradients. It holds for linear VJPs of correct parent shapes, in particular for every tape produced by lowerGraphToTape.

          Instances For

            A successful getNode? witnesses that the id is in range.

            Conversely, any in-range id resolves to a node. The pair lets the invariants below be phrased in terms of getNode? alone, without carrying array bounds around.

            Totalizing a well-formed optional accumulator yields a well-formed total accumulator.

            theorem Proofs.Autograd.Algebra.Graph.optGradsOk_set {α : Type} [TorchLean.Storage α] {t : Runtime.Autograd.Tape (Spec.SomeTensor α)} {grads : Array (Option (Spec.SomeTensor α))} (hok : OptGradsOk t grads) {id : ℕ} {node : Runtime.Autograd.Node (Spec.SomeTensor α)} (hnode : t.getNode? id = some node) {g : Spec.SomeTensor α} (hg : g.shape = node.value.shape) (hid : id < grads.size) :
            OptGradsOk t (grads.set id (some g) hid)

            Writing a correctly shaped gradient preserves OptGradsOk.

            Adding a zero contribution is a no-op #

            Adding the all-zero tensor on the left is the identity.

            Adding the all-zero tensor on the right is the identity.

            Accumulating onto a node's zero cotangent returns the contribution unchanged.

            Accumulating a node's zero cotangent onto an existing gradient leaves it unchanged.

            addGradAll with a parent's zero cotangent is the identity on a well-formed accumulator.

            theorem Proofs.Autograd.Algebra.Graph.foldlM_addGradAll_zero {α : Type} [TorchLean.Storage α] [AddZeroClass α] {t : Runtime.Autograd.Tape (Spec.SomeTensor α)} {grads : Array (Spec.SomeTensor α)} (hok : TotGradsOk t grads) (contribs : List (ℕ × Spec.SomeTensor α)) (hz : ∀ (pid : ℕ) (pg : Spec.SomeTensor α), (pid, pg) ∈ contribs → ∃ (pnode : Runtime.Autograd.Node (Spec.SomeTensor α)), t.getNode? pid = some pnode ∧ pg = zeroCotangent pnode) :
            List.foldlM (fun (acc2 : Array (Spec.SomeTensor α)) (x : ℕ × Spec.SomeTensor α) => match x with | (pid, pg) => t.addGradAll acc2 pid pg) grads contribs = Except.ok grads

            Folding zero contributions through addGradAll is the identity on a well-formed accumulator.

            One contribution: addGradDense simulates addGradAll #

            theorem Proofs.Autograd.Algebra.Graph.addGradDense_optGradsOk {α : Type} [TorchLean.Storage α] [Add α] {t : Runtime.Autograd.Tape (Spec.SomeTensor α)} {grads : Array (Option (Spec.SomeTensor α))} (hok : OptGradsOk t grads) {pid : ℕ} {g : Spec.SomeTensor α} {grads' : Array (Option (Spec.SomeTensor α))} (h : t.addGradDense grads pid g = Except.ok grads') :
            OptGradsOk t grads'

            addGradDense preserves the optional accumulator invariant.

            One addGradDense step, mapped through Tape.totalize, is the matching addGradAll step.

            Folding contributions #

            theorem Proofs.Autograd.Algebra.Graph.except_map_bind {β γ : Type} (f : β → γ) {x : Runtime.Autograd.Result β} {y : Runtime.Autograd.Result γ} {k : β → Runtime.Autograd.Result β} {k' : γ → Runtime.Autograd.Result γ} (hx : Except.map f x = y) (hk : ∀ (b : β), x = Except.ok b → Except.map f (k b) = k' (f b)) :
            Except.map f (x >>= k) = y >>= k'

            Except.map commutes with bind when it commutes with the first action and, on success, with the continuation. This is the single lemma behind every "executed step simulates proved step" composition below.

            theorem Proofs.Autograd.Algebra.Graph.foldlM_addGradDense_optGradsOk {α : Type} [TorchLean.Storage α] [Add α] {t : Runtime.Autograd.Tape (Spec.SomeTensor α)} (contribs : List (ℕ × Spec.SomeTensor α)) {grads grads' : Array (Option (Spec.SomeTensor α))} :
            OptGradsOk t grads → List.foldlM (fun (acc2 : Array (Option (Spec.SomeTensor α))) (x : ℕ × Spec.SomeTensor α) => match x with | (pid, pg) => t.addGradDense acc2 pid pg) grads contribs = Except.ok grads' → OptGradsOk t grads'

            Folding addGradDense preserves the optional accumulator invariant.

            theorem Proofs.Autograd.Algebra.Graph.foldlM_addGradDense_map_totalize {α : Type} [TorchLean.Storage α] [AddZeroClass α] {t : Runtime.Autograd.Tape (Spec.SomeTensor α)} (contribs : List (ℕ × Spec.SomeTensor α)) {grads : Array (Option (Spec.SomeTensor α))} :
            OptGradsOk t grads → Except.map t.totalize (List.foldlM (fun (acc2 : Array (Option (Spec.SomeTensor α))) (x : ℕ × Spec.SomeTensor α) => match x with | (pid, pg) => t.addGradDense acc2 pid pg) grads contribs) = List.foldlM (fun (acc2 : Array (Spec.SomeTensor α)) (x : ℕ × Spec.SomeTensor α) => match x with | (pid, pg) => t.addGradAll acc2 pid pg) (t.totalize grads) contribs

            Folding addGradDense, mapped through Tape.totalize, is folding addGradAll.

            One node: the executed step simulates the proved step #

            theorem Proofs.Autograd.Algebra.Graph.backwardDenseStep_optGradsOk {α : Type} [TorchLean.Storage α] [Add α] {t : Runtime.Autograd.Tape (Spec.SomeTensor α)} {grads grads' : Array (Option (Spec.SomeTensor α))} (hok : OptGradsOk t grads) {id : ℕ} (h : t.backwardDenseStep grads id = Except.ok grads') :
            OptGradsOk t grads'

            The executed step preserves the optional accumulator invariant.

            The executed step at a node, mapped through Tape.totalize, is the proved step backwardDenseFromStep at the same node, provided the tape is zero preserving.

            When the node was reached both sides run the same VJP on the same cotangent. When it was not, the executed side does nothing and the proved side runs the VJP on the zero cotangent, whose contributions are all zeros and therefore accumulate to nothing.

            The reverse loop #

            theorem Proofs.Autograd.Algebra.Graph.foldlM_range_reverse_succ {m : Type → Type} [Monad m] {β : Type} (f : β → ℕ → m β) (b : β) (n : ℕ) :
            List.foldlM f b (List.range (n + 1)).reverse = do let b' ← f b n List.foldlM f b' (List.range n).reverse

            Peel the first (largest) id off a reverse-range fold.

            The executed reverse loop over the first n node ids, mapped through Tape.totalize, is the proved loop backwardDenseFromLoop over the same ids.

            Main theorem #

            theorem Proofs.Autograd.Algebra.Graph.optGradsOk_seed {α : Type} [TorchLean.Storage α] {t : Runtime.Autograd.Tape (Spec.SomeTensor α)} {outId : ℕ} {outNode : Runtime.Autograd.Node (Spec.SomeTensor α)} (hout : t.getNode? outId = some outNode) {seed : Spec.SomeTensor α} (hseed : seed.shape = outNode.value.shape) (hlt : outId < t.nodes.size) :
            OptGradsOk t ((Array.replicate t.nodes.size none).set outId (some seed) ⋯)

            The seeded optional accumulator backwardDense starts from.

            theorem Proofs.Autograd.Algebra.Graph.totalize_seed {α : Type} [TorchLean.Storage α] [Zero α] (t : Runtime.Autograd.Tape (Spec.SomeTensor α)) (outId : ℕ) (seed : Spec.SomeTensor α) (hlt : outId < t.nodes.size) :
            t.totalize ((Array.replicate t.nodes.size none).set outId (some seed) ⋯) = oneHotGrads t outId seed

            Totalizing the seeded optional accumulator gives the one-hot total accumulator.

            backwardDense is the reverse fold of backwardDenseStep from the seeded accumulator.

            Executed backward pass = proved backward pass. On a zero-preserving tape, the totalized executed sweep backwardDenseAll (which skips unreached nodes) returns exactly what the proved sweep backwardDenseFrom returns when started from the one-hot seed array, including agreement of the error cases. The hypotheses on outId and seed are the checks backwardDense performs before traversing; without them backwardDense fails while backwardDenseFrom may not.