TorchLean API

NN.Proofs.Autograd.Runtime.Link.Core

Link #

Link the executable runtime tape (Runtime.Autograd.Tape) to the shape-indexed SSA/DAG models in Proofs.Autograd.Algebra.

GraphData stores executable forward, JVP, and VJP functions without derivative laws. Graph extends that representation with the local adjointness law used by its global backpropagation theorem. Lowering places the stored VJP into each runtime node's backward closure.

What is proved here #

The core invariant making the runtime reverse loop well-founded is that lowered nodes only emit contributions to earlier node ids (pid < id).

PyTorch correspondence / citations #

This is analogous to lowering a graph representation to an executable autograd tape whose nodes carry backward closures (PyTorch does this internally for the eager autograd engine). https://pytorch.org/docs/stable/autograd.html

Extend a tape with leaf nodes for every tensor in the input context Γ.

Each leaf has requiresGrad = true and an empty backward contribution array, so the runtime loop treats them as gradient accumulation slots but never produces parent contributions from them.

Instances For

    Turn a shape-erased tensor into a runtime leaf node.

    This is the node-level counterpart of addLeaves: it has no parents and contributes nothing in backward.

    Instances For

      Result equations #

      Result is Except String. The runtime backward functions are written with do notation, so proofs about them repeatedly need the unfolding equations for bind, pure, throw and Except.map. They are stated once here and used with simp only instead of unfolding the Monad instances by hand.

      theorem Proofs.Autograd.Algebra.Graph.result_bind_ok {β γ : Type} (a : β) (f : βRuntime.Autograd.Result γ) :
      Except.ok a >>= f = f a

      Binding a successful Result applies the continuation.

      Binding a failed Result propagates the error.

      theorem Proofs.Autograd.Algebra.Graph.result_map_ok {β γ : Type} (f : βγ) (a : β) :

      Except.map on a successful Result.

      Except.map on a failed Result.

      Backward closures of t only point backwards: every contribution emitted by the node stored at id targets a node id strictly smaller than id. This is the invariant that makes the dense reverse loop well-founded; lowerGraphToTape_backward_pids_lt_id and its GraphData counterpart establish it for lowered tapes.

      Instances For

        Each entry of toShapeErasedArray carries the shape recorded at its position of ss.

        addLeaves grows the tape by exactly Γ.length nodes.

        Value projection of nodes_addLeaves: node.value agrees with toShapeErasedArray for added leaves.

        def Proofs.Autograd.Algebra.Graph.lowerNode {α Δ : Type} [TorchLean.Storage α] {Γ : List Spec.Shape} {τ : Spec.Shape} (name : Option String) (node : NodeData α Δ Γ τ) (ctx : TorchLean.TensorPack α Γ) (d : Δ) :

        Runtime node produced by lowering one graph node.

        The node stores the forward value and wraps the stored vjp in a runtime backward closure that checks the shape of the upstream cotangent and emits one contribution per entry of the node's input context, indexed from 0. The name only serves debugging output.

        Instances For
          @[simp]
          theorem Proofs.Autograd.Algebra.Graph.lowerNode_value {α Δ : Type} [TorchLean.Storage α] {Γ : List Spec.Shape} {τ : Spec.Shape} (name : Option String) (node : NodeData α Δ Γ τ) (ctx : TorchLean.TensorPack α Γ) (d : Δ) :
          (lowerNode name node ctx d).value = Spec.SomeTensor.ofTensor (node.forward ctx d)

          The lowered node stores the node's forward value.

          @[simp]
          theorem Proofs.Autograd.Algebra.Graph.lowerNode_requiresGrad {α Δ : Type} [TorchLean.Storage α] {Γ : List Spec.Shape} {τ : Spec.Shape} (name : Option String) (node : NodeData α Δ Γ τ) (ctx : TorchLean.TensorPack α Γ) (d : Δ) :
          (lowerNode name node ctx d).requiresGrad = true

          Lowered nodes always take part in gradient accumulation.

          theorem Proofs.Autograd.Algebra.Graph.lowerNode_backward_of_shape {α Δ : Type} [TorchLean.Storage α] {Γ : List Spec.Shape} {τ : Spec.Shape} (name : Option String) (node : NodeData α Δ Γ τ) (ctx : TorchLean.TensorPack α Γ) (d : Δ) (v : Spec.SomeTensor α) (h : v.shape = τ) :
          (lowerNode name node ctx d).backward v = Except.ok ((node.vjp ctx d (v.cast h)).toIndexedShapeErasedArray 0)

          On an upstream cotangent of the right shape, the lowered backward runs the stored vjp.

          theorem Proofs.Autograd.Algebra.Graph.lowerNode_backward_of_ne {α Δ : Type} [TorchLean.Storage α] {Γ : List Spec.Shape} {τ : Spec.Shape} (name : Option String) (node : NodeData α Δ Γ τ) (ctx : TorchLean.TensorPack α Γ) (d : Δ) (v : Spec.SomeTensor α) (h : v.shape τ) :
          (lowerNode name node ctx d).backward v = Except.error "autograd: upstream gradient shape mismatch"

          On an upstream cotangent of the wrong shape, the lowered backward fails.

          Lower an executable graph (GraphData) to a runtime tape by evaluating forward nodes and storing each node's vjp program in its runtime backward closure.

          PyTorch analogy: this corresponds to building a tape of autograd nodes during the forward pass, where each node stores enough information to compute parent contributions when given an upstream cotangent.

          Instances For

            Forward-pass correspondence #

            The next lemmas show that lowerGraphDataToTape preserves executable forward semantics, and that the resulting runtime tape contains exactly the evaluated context as shape-erased tensors in order.

            The context returned by lowerGraphDataToTape agrees with GraphData.eval.

            The lowered tape's .value array is GraphData.eval with shapes erased in the same order.

            Size bookkeeping: the lowered tape contains one runtime node for each element of Γ ++ ss.

            Lower a proved graph (Graph) to a runtime tape by evaluating forward nodes and storing each node’s proved vjp.

            Compared to lowerGraphDataToTape, this uses the pure graph interface (no explicit GraphData payload).

            Instances For
              theorem Proofs.Autograd.Algebra.Graph.lowerGraphToTape_ctx_eq_eval {α Δ : Type} [TorchLean.Storage α] [CommSemiring α] {Γ ss : List Spec.Shape} (g : Graph Δ Γ ss) (x : TorchLean.TensorPack α Γ) (d : Δ) :
              (g.lowerGraphToTape x d).2 = g.eval x d

              The context returned by lowerGraphToTape agrees with the proved Graph.eval.

              The lowered tape's .value array is Graph.eval with shapes erased in the same order.

              Size bookkeeping: lowerGraphToTape produces Γ.length + ss.length runtime nodes.

              Full backpropagation (dense) for proofs and runtime #

              The runtime engine computes a dense gradient array, accumulating cotangents for every node in the tape (inputs and intermediates). The following definition and theorems connect that behavior to the proved backpropagation semantics.

              def Proofs.Autograd.Algebra.Graph.backpropAllCtx {α Δ : Type} [TorchLean.Storage α] [CommSemiring α] {Γ ss : List Spec.Shape} (g : Graph Δ Γ ss) (x : TorchLean.TensorPack α Γ) (d : Δ) (seed : TorchLean.TensorPack α (Γ ++ ss)) :

              A "full" backpropagation that returns gradients for every value in Γ ++ ss.

              Instances For
                def Proofs.Autograd.Algebra.GraphData.backpropAllCtx {α Δ : Type} [TorchLean.Storage α] [Add α] {Γ ss : List Spec.Shape} (g : GraphData α Δ Γ ss) (x : TorchLean.TensorPack α Γ) (d : Δ) (seed : TorchLean.TensorPack α (Γ ++ ss)) :

                “Full” backpropagation for GraphData that returns gradients for every value in Γ ++ ss, including inputs.

                This is the GraphData-analogue of backpropAllCtx above. We keep both definitions because:

                • Graph uses [CommSemiring α] (so it can express dot products and semiring-based accumulation), while
                • GraphData only needs [Add α] here (it just adds contributions).

                Both follow the same reverse-mode accumulation structure: peel off the last node, apply its VJP to the seed on that node, add into the previous seed, and recurse.

                Instances For