Reading intermediate values out of an evaluated tape graph

Graph.evalVec returns the whole flattened context, inputs followed by every saved intermediate. The pointwise normalization proofs need to know what a specific block of that vector is: the variance block of the LayerNorm prefix, for instance, must be shown nonnegative before sqrt and inv can be differentiated.

This file gives the coordinate description of CtxVec.get and the three rules that let a proof walk back through a snoc chain:

• Graph.get_evalVec_snoc_last: the block just appended is the node's forward value;
• Graph.get_evalVec_snoc_of_lt: any earlier block is unchanged by appending a node;
• Graph.get_evalVec_input: an input block is unchanged by the whole graph.

theorem Proofs.Autograd.Idx.ext' {Γ : List Spec.Shape} {s : Spec.Shape} {a b : Idx Γ s} (h : a.i = b.i) : a = b

Two typed indices with the same position are equal.

def Proofs.Autograd.CtxVec.blockOffset : List Spec.Shape → ℕ → ℕ

Offset of the block at list position k inside the flattened context vector.

theorem Proofs.Autograd.CtxVec.blockOffset_add_lt {Γ : List Spec.Shape} (i : Fin Γ.length) (j : Fin (Γ.get i).size) : blockOffset Γ ↑i + ↑j < ctxSize Γ

Every coordinate of a block lies inside the flattened context.

theorem Proofs.Autograd.CtxVec.getBlock_ofLp {Γ : List Spec.Shape} (i : Fin Γ.length) (x : CtxVec Γ) (j : Fin (Γ.get i).size) : (getBlock i x).ofLp j = x.ofLp ⟨blockOffset Γ ↑i + ↑j, ⋯⟩

Coordinates of getBlock are coordinates of the context, shifted by the block offset.

theorem Proofs.Autograd.CtxVec.blockOffset_add_lt' {Γ : List Spec.Shape} {s : Spec.Shape} (idx : Idx Γ s) (j : Fin s.size) : blockOffset Γ ↑idx.i + ↑j < ctxSize Γ

Every coordinate of the block selected by a typed index lies inside the context.

theorem Proofs.Autograd.CtxVec.get_ofLp {Γ : List Spec.Shape} {s : Spec.Shape} (idx : Idx Γ s) (x : CtxVec Γ) (j : Fin s.size) : (get idx x).ofLp j = x.ofLp ⟨blockOffset Γ ↑idx.i + ↑j, ⋯⟩

Coordinate description of CtxVec.get.

theorem Proofs.Autograd.CtxVec.blockOffset_append_length (Γ rest : List Spec.Shape) : blockOffset (Γ ++ rest) Γ.length = ctxSize Γ

The block right after a prefix Γ starts at ctxSize Γ.

theorem Proofs.Autograd.CtxVec.blockOffset_append_of_lt (Γ rest : List Spec.Shape) (k : ℕ) : k < Γ.length → blockOffset (Γ ++ rest) k = blockOffset Γ k

Appending shapes does not move the blocks of the prefix.

theorem Proofs.Autograd.CtxVec.getBlock_flattenCtx_zero {s : Spec.Shape} {ss : List Spec.Shape} (X : TorchLean.Tensor ℝ s) (xs : TorchLean.TensorPack ℝ ss) (h : 0 < (s :: ss).length) : getBlock ⟨0, h⟩ (flattenCtx (TorchLean.TensorPack.cons X xs)) = tensorToVec X

The first block of a flattened tensor pack is the first tensor.

theorem Proofs.Autograd.CtxVec.getBlock_flattenCtx_succ {s : Spec.Shape} {ss : List Spec.Shape} (X : TorchLean.Tensor ℝ s) (xs : TorchLean.TensorPack ℝ ss) (k : ℕ) (hk : k + 1 < (s :: ss).length) : getBlock ⟨k + 1, hk⟩ (flattenCtx (TorchLean.TensorPack.cons X xs)) = getBlock ⟨k, ⋯⟩ (flattenCtx xs)

Later blocks of a flattened tensor pack are blocks of the tail.

theorem Proofs.Autograd.castCtxVec_ofLp {Γ₁ Γ₂ : List Spec.Shape} (h : Γ₁ = Γ₂) (x : CtxVec Γ₁) (k : Fin (ctxSize Γ₂)) : (castCtxVec h x).ofLp k = x.ofLp (Fin.cast ⋯ k)

Coordinates of castCtxVec are coordinates of the original vector.

theorem Proofs.Autograd.snocCtx_ofLp_of_lt {Γ : List Spec.Shape} {τ : Spec.Shape} (ctx : CtxVec Γ) (t : Vec τ.size) (k : Fin (ctxSize (Γ ++ [τ]))) (hk : ↑k < ctxSize Γ) : (snocCtx ctx t).ofLp k = ctx.ofLp ⟨↑k, hk⟩

Coordinates of snocCtx below the prefix size read the prefix.

theorem Proofs.Autograd.snocCtx_ofLp_last {Γ : List Spec.Shape} {τ : Spec.Shape} (ctx : CtxVec Γ) (t : Vec τ.size) (k : Fin (ctxSize (Γ ++ [τ]))) (j : Fin τ.size) (hk : ↑k = ctxSize Γ + ↑j) : (snocCtx ctx t).ofLp k = t.ofLp j

Coordinates of snocCtx at or above the prefix size read the appended block.

theorem Proofs.Autograd.Graph.get_evalVec_nil {Γ : List Spec.Shape} {s : Spec.Shape} (xV : CtxVec Γ) (idx : Idx (Γ ++ []) s) (idx' : Idx Γ s) (hv : ↑idx.i = ↑idx'.i) : CtxVec.get idx (nil.evalVec xV) = CtxVec.get idx' xV

Evaluating the empty graph leaves every block in place.

theorem Proofs.Autograd.Graph.get_evalVec_snoc_last {Γ ss : List Spec.Shape} {τ : Spec.Shape} (g : Graph Γ ss) (node : Node (Γ ++ ss) τ) (xV : CtxVec Γ) (idx : Idx (Γ ++ (ss ++ [τ])) τ) (hv : ↑idx.i = (Γ ++ ss).length) : CtxVec.get idx ((g.snoc node).evalVec xV) = node.forwardVec (g.evalVec xV)

The block appended by snoc is the node's forward value on the prefix evaluation.

theorem Proofs.Autograd.Graph.get_evalVec_snoc_of_lt {Γ ss : List Spec.Shape} {τ s : Spec.Shape} (g : Graph Γ ss) (node : Node (Γ ++ ss) τ) (xV : CtxVec Γ) (idx : Idx (Γ ++ (ss ++ [τ])) s) (idx' : Idx (Γ ++ ss) s) (hv : ↑idx.i = ↑idx'.i) : CtxVec.get idx ((g.snoc node).evalVec xV) = CtxVec.get idx' (g.evalVec xV)

Blocks below the appended one are unchanged by snoc.

theorem Proofs.Autograd.Graph.get_evalVec_weaken {Γ : List Spec.Shape} {s : Spec.Shape} (idx : Idx Γ s) {ss : List Spec.Shape} (g : Graph Γ ss) (xV : CtxVec Γ) : CtxVec.get (idx.weaken ss) (g.evalVec xV) = CtxVec.get idx xV

Input blocks are unchanged by evaluating any graph.

theorem Proofs.Autograd.Graph.get_evalVec_input {Γ ss : List Spec.Shape} {s : Spec.Shape} (g : Graph Γ ss) (xV : CtxVec Γ) (idx : Idx (Γ ++ ss) s) (idx' : Idx Γ s) (hv : ↑idx.i = ↑idx'.i) : CtxVec.get idx (g.evalVec xV) = CtxVec.get idx' xV

Input blocks are unchanged by evaluating any graph, for any index into the input prefix.