Differentiable graph composition

The DGraph wrapper packages a tape graph together with node-local NodeFDerivCorrect proofs. A graph is built incrementally with DGraph.snoc, DGraph.append, and DGraph.weakenContext, after which DGraph.backpropVec_eq_adjoint_fderiv applies without threading a separate proof object. These adapters let model-level VJP theorems reuse the proved primitive nodes rather than reproving backprop correctness for each architecture from scratch.

noncomputable def Proofs.Autograd.TapeNodes.takeLeftCLM (m n : ℕ) : Vec (m + n) →L[ℝ] Vec m

takeLeftVec packaged as a continuous linear map.

noncomputable def Proofs.Autograd.TapeNodes.takeRightCLM (m n : ℕ) : Vec (m + n) →L[ℝ] Vec n

takeRightVec packaged as a continuous linear map.

structure Proofs.Autograd.DGraph (Γ ss : List Spec.Shape) : Type

A tape graph bundled with the proof that every one of its nodes is analytically correct.

• g : Graph Γ ss

  The underlying tape/DAG graph.

• hg : GraphFDerivCorrect self.g

  Proof that every node is analytically correct (jvp = fderiv).

noncomputable def Proofs.Autograd.DGraph.dropMiddleCLM (Γ extra ss : List Spec.Shape) : CtxVec (Γ ++ extra ++ ss) →L[ℝ] CtxVec (Γ ++ ss)

Drop an unused middle context block from a graph node context.

When a graph dg : DGraph Γ ss is reused inside a larger context Γ ++ extra, a node that originally reads Γ ++ ss is evaluated in the actual context (Γ ++ extra) ++ ss. This projection keeps the original inputs Γ and the already-computed intermediates ss, and ignores the carried parameters in extra.

@[simp]

theorem Proofs.Autograd.DGraph.dropMiddleCLM_apply {Γ extra ss : List Spec.Shape} (x : CtxVec (Γ ++ extra ++ ss)) : (dropMiddleCLM Γ extra ss) x = castVec ⋯ (appendVec (TapeNodes.takeLeftVec (castVec ⋯ (TapeNodes.takeLeftVec (castVec ⋯ x)))) (TapeNodes.takeRightVec (castVec ⋯ x)))

Unfolds the middle-dropping projection into its take, append and cast components.

Composing two tapes leaves the inner tape's private saved values in the middle of the context; this map is what discards them, keeping the caller's context and the composed outputs.

noncomputable def Proofs.Autograd.DGraph.weakenNodeMiddle {Γ extra ss : List Spec.Shape} {τ : Spec.Shape} (node : Node (Γ ++ ss) τ) : Node (Γ ++ extra ++ ss) τ

Reuse a node in a context that carries extra unused inputs between the original inputs and the current SSA intermediates.

The VJP is obtained by applying the adjoint of dropMiddleCLM, so gradients land only in the original inputs and previous intermediates; the extra carried parameters receive zero contribution from this reused node.

noncomputable def Proofs.Autograd.DGraph.weakenNodeMiddleFDerivCorrect {Γ extra ss : List Spec.Shape} {τ : Spec.Shape} {node : Node (Γ ++ ss) τ} (hn : NodeFDerivCorrect node) : NodeFDerivCorrect (weakenNodeMiddle node)

Transport a global node derivative certificate across weakenNodeMiddle.

def Proofs.Autograd.DGraph.nil {Γ : List Spec.Shape} : DGraph Γ []

Empty differentiable graph.

def Proofs.Autograd.DGraph.snoc {Γ ss : List Spec.Shape} {τ : Spec.Shape} (dg : DGraph Γ ss) (node : Node (Γ ++ ss) τ) (hn : NodeFDerivCorrect node) : DGraph Γ (ss ++ [τ])

Append a node together with its NodeFDerivCorrect certificate.

def Proofs.Autograd.DGraph.castNodeContext {Γ₁ Γ₂ : List Spec.Shape} {τ : Spec.Shape} (h : Γ₁ = Γ₂) (node : Node Γ₁ τ) : Node Γ₂ τ

Transport a node across a definitional/context-list equality.

This is mostly used by graph composition: the second graph sees its context as (Γ ++ ss₁) ++ ss₂, while the composed graph sees the same values as Γ ++ (ss₁ ++ ss₂).

def Proofs.Autograd.DGraph.castNodeFDerivCorrect {Γ₁ Γ₂ : List Spec.Shape} {τ : Spec.Shape} (h : Γ₁ = Γ₂) {node : Node Γ₁ τ} (hn : NodeFDerivCorrect node) : NodeFDerivCorrect (castNodeContext h node)

Transport a node F-derivative certificate along castNodeContext.

noncomputable def Proofs.Autograd.DGraph.Internal.append {Γ ss₁ ss₂ : List Spec.Shape} (dg₁ : DGraph Γ ss₁) (g₂ : Graph (Γ ++ ss₁) ss₂) (hg₂ : GraphFDerivCorrect g₂) : DGraph Γ (ss₁ ++ ss₂)

Recursive implementation for DGraph.append, stated over an explicit graph and proof.

noncomputable def Proofs.Autograd.DGraph.Internal.weakenContext {Γ ss : List Spec.Shape} (extra : List Spec.Shape) (g : Graph Γ ss) (hg : GraphFDerivCorrect g) : DGraph (Γ ++ extra) ss

Recursive implementation for DGraph.weakenContext, stated over an explicit graph and proof.

noncomputable def Proofs.Autograd.DGraph.append {Γ ss₁ ss₂ : List Spec.Shape} (dg₁ : DGraph Γ ss₁) (dg₂ : DGraph (Γ ++ ss₁) ss₂) : DGraph Γ (ss₁ ++ ss₂)

Append a proof-carrying graph after another proof-carrying graph.

If dg₁ : DGraph Γ ss₁ has already computed some SSA values, then a second graph dg₂ : DGraph (Γ ++ ss₁) ss₂ may use both the original inputs and those saved values. append turns the pair into one DGraph Γ (ss₁ ++ ss₂).

This is the general composition adapter needed for model-level proofs: residual attention can feed LayerNorm, a recurrent cell can feed the next unrolled step, and larger blocks can be assembled while reusing the existing node-level correctness proofs.

noncomputable def Proofs.Autograd.DGraph.weakenContext {Γ ss : List Spec.Shape} (dg : DGraph Γ ss) (extra : List Spec.Shape) : DGraph (Γ ++ extra) ss

Run a proof-carrying graph while carrying extra unused inputs.

If dg : DGraph Γ ss, then weakenContext dg extra : DGraph (Γ ++ extra) ss evaluates the same nodes while preserving an enlarged input context. Each reused node sees the projection Γ ++ ss_so_far of the actual context (Γ ++ extra) ++ ss_so_far; gradients are inserted back by the adjoint projection, so the carried extras receive no gradient contribution from nodes that do not read them.

theorem Proofs.Autograd.DGraph.backpropVec_eq_adjoint_fderiv {Γ ss : List Spec.Shape} (dg : DGraph Γ ss) (xV : CtxVec Γ) (seedV : CtxVec (Γ ++ ss)) : dg.g.backpropVec xV seedV = (ContinuousLinearMap.adjoint (fderiv ℝ dg.g.evalVec xV)) seedV

End-to-end analytic theorem for bundled graphs.

This is just Graph.backpropVec_eq_adjoint_fderiv with the bundled proof dg.hg.