Reduction and shape tape nodes

Scalar sums, broadcast-to, reduce-sum, reduce-mean, concatenation, and the linear shape adapters used by larger graph proofs.

noncomputable def Proofs.Autograd.TapeNodes.vecScalarCLM : ℝ →L[ℝ] Vec Spec.Shape.scalar.size

Continuous linear map embedding a scalar into the 1D scalar-vector representation.

@[simp]

theorem Proofs.Autograd.TapeNodes.vecScalarCLM_apply (a : ℝ) (i : Fin Spec.Shape.scalar.size) : (vecScalarCLM a).ofLp i = a

The one coordinate of a scalar vector is the scalar.

@[simp]

theorem Proofs.Autograd.TapeNodes.vecScalarCLM_ofLp (a : ℝ) (i : Fin Spec.Shape.scalar.size) : (vecScalarCLM a).ofLp i = a

Same through the WithLp wrapper.

def Proofs.Autograd.TapeNodes.sumCLM (n : ℕ) : Vec n →L[ℝ] ℝ

Continuous linear map summing the entries of a vector: x ↦ ∑ i, x i.

theorem Proofs.Autograd.TapeNodes.sumCLM_apply {n : ℕ} (x : Vec n) : (sumCLM n) x = ∑ i : Fin n, x.ofLp i

Evaluation lemma for sumCLM.

noncomputable def Proofs.Autograd.TapeNodes.sum {Γ : List Spec.Shape} {s : Spec.Shape} (idx : Idx Γ s) : Node Γ Spec.Shape.scalar

Sum all entries of a context tensor into a scalar tensor.

noncomputable def Proofs.Autograd.TapeNodes.sumFderiv {Γ : List Spec.Shape} {s : Spec.Shape} (idx : Idx Γ s) : NodeFDerivCorrect (sum idx)

NodeFDerivCorrect for sum: derivative is the composite of context projection and coordinate sum.

def Proofs.Autograd.TapeNodes.Broadcast.broadcastToIndex {s₁ s₂ : Spec.Shape} : s₁.CanBroadcastTo s₂ → Fin s₂.size → Fin s₁.size

Compute the source index in s₁ that corresponds to a target index in s₂ under broadcasting.

The recursion is on the shapes. The broadcast relation only rules out the impossible cases and shows that a source extent differing from the target extent is one.

noncomputable def Proofs.Autograd.TapeNodes.Broadcast.broadcastToVec {s₁ s₂ : Spec.Shape} (cb : s₁.CanBroadcastTo s₂) : Vec s₁.size → Vec s₂.size

Broadcast a vector Vec (size s₁) into Vec (size s₂) using the CanBroadcastTo index map.

noncomputable def Proofs.Autograd.TapeNodes.Broadcast.broadcastToCLM {s₁ s₂ : Spec.Shape} (cb : s₁.CanBroadcastTo s₂) : Vec s₁.size →L[ℝ] Vec s₂.size

Continuous-linear-map form of broadcastToVec.

@[simp]

theorem Proofs.Autograd.TapeNodes.Broadcast.broadcastToCLM_apply {s₁ s₂ : Spec.Shape} (cb : s₁.CanBroadcastTo s₂) (v : Vec s₁.size) : (broadcastToCLM cb) v = broadcastToVec cb v

The bundled broadcast map computes broadcastToVec. Broadcasting is linear, so nothing about the shape relation cb needs to reappear in the derivative.

def Proofs.Autograd.TapeNodes.Broadcast.afterSumIndex (s : Spec.Shape) (axis : ℕ) : Fin s.size → Fin (TorchLean.Tensor.shapeAfterSum s axis).size

Source index obtained by deleting coordinate axis from an index into s.

noncomputable def Proofs.Autograd.TapeNodes.Broadcast.afterSumVec (s : Spec.Shape) (axis : ℕ) : Vec (TorchLean.Tensor.shapeAfterSum s axis).size → Vec s.size

Repeat a flattened reduced tensor along the axis removed by shapeAfterSum.

noncomputable def Proofs.Autograd.TapeNodes.Broadcast.afterSumCLM (s : Spec.Shape) (axis : ℕ) : Vec (TorchLean.Tensor.shapeAfterSum s axis).size →L[ℝ] Vec s.size

Continuous-linear-map form of afterSumVec.

@[simp]

theorem Proofs.Autograd.TapeNodes.Broadcast.afterSumCLM_apply (s : Spec.Shape) (axis : ℕ) (v : Vec (TorchLean.Tensor.shapeAfterSum s axis).size) : (afterSumCLM s axis) v = afterSumVec s axis v

The bundled repeat-along-the-summed-axis map computes afterSumVec.

This map is the adjoint of summation over axis, which is why the backward pass for a reduction is a broadcast: each output coordinate contributed to exactly one sum.

noncomputable def Proofs.Autograd.TapeNodes.broadcastTo {Γ : List Spec.Shape} {s₁ s₂ : Spec.Shape} (idx : Idx Γ s₁) (cb : s₁.CanBroadcastTo s₂) : Node Γ s₂

General shape broadcast node s₁ → s₂ (linear).

noncomputable def Proofs.Autograd.TapeNodes.broadcastToFderiv {Γ : List Spec.Shape} {s₁ s₂ : Spec.Shape} (idx : Idx Γ s₁) (cb : s₁.CanBroadcastTo s₂) : NodeFDerivCorrect (broadcastTo idx cb)

NodeFDerivCorrect for broadcastTo (broadcasting is linear).

noncomputable def Proofs.Autograd.TapeNodes.reduceSum {Γ : List Spec.Shape} {s : Spec.Shape} (axis : ℕ) [_valid : Spec.Shape.HasNonemptyAxis axis s] [_wf : s.WellFormed] (idx : Idx Γ s) : Node Γ (TorchLean.Tensor.shapeAfterSum s axis)

Sum reduction along axis (linear; adjoint is broadcast back).

noncomputable def Proofs.Autograd.TapeNodes.reduceSumFderiv {Γ : List Spec.Shape} {s : Spec.Shape} (axis : ℕ) [valid : Spec.Shape.HasNonemptyAxis axis s] [wf : s.WellFormed] (idx : Idx Γ s) : NodeFDerivCorrect (reduceSum axis idx)

NodeFDerivCorrect for reduceSum.

noncomputable def Proofs.Autograd.TapeNodes.reduceMean {Γ : List Spec.Shape} {s : Spec.Shape} (axis : ℕ) [valid : Spec.Shape.HasNonemptyAxis axis s] [_wf : s.WellFormed] (idx : Idx Γ s) : Node Γ (TorchLean.Tensor.shapeAfterSum s axis)

Mean reduction along axis (linear; adjoint is broadcast+scale).

noncomputable def Proofs.Autograd.TapeNodes.reduceMeanFderiv {Γ : List Spec.Shape} {s : Spec.Shape} (axis : ℕ) [valid : Spec.Shape.HasNonemptyAxis axis s] [wf : s.WellFormed] (idx : Idx Γ s) : NodeFDerivCorrect (reduceMean axis idx)

NodeFDerivCorrect for reduceMean.

noncomputable def Proofs.Autograd.TapeNodes.takeLeftVec {m n : ℕ} (v : Vec (m + n)) : Vec m

Take the left m entries of a length m+n vector.

noncomputable def Proofs.Autograd.TapeNodes.takeRightVec {m n : ℕ} (v : Vec (m + n)) : Vec n

Take the right n entries of a length m+n vector.

noncomputable def Proofs.Autograd.TapeNodes.concatLeadingAxis {Γ : List Spec.Shape} {n m : ℕ} {s : Spec.Shape} (a : Idx Γ (Spec.Shape.dim n s)) (b : Idx Γ (Spec.Shape.dim m s)) : Node Γ (Spec.Shape.dim (n + m) s)

Concatenate two tensors along dimension 0 (dim-0 concat), using flattened vectors internally.

noncomputable def Proofs.Autograd.TapeNodes.concatLeadingAxisFderiv {Γ : List Spec.Shape} {n m : ℕ} {s : Spec.Shape} (a : Idx Γ (Spec.Shape.dim n s)) (b : Idx Γ (Spec.Shape.dim m s)) : NodeFDerivCorrect (concatLeadingAxis a b)

NodeFDerivCorrect for concatLeadingAxis (concat is linear).