Shape

Additional analytic (HasFDerivAt) tape nodes for shape permutations.

These nodes are linear/isometric and are useful for models that do explicit reshaping and dimension permutations (e.g. Multi-Head Attention head splitting/combining).

theorem Proofs.Autograd.TapeNodes.ShapeOps.inner_castVec_left {n m : ℕ} (h : n = m) (x : Vec n) (y : Vec m) : inner ℝ (castVec h x) y = inner ℝ x (castVec ⋯ y)

Move castVec across the left argument of an inner product.

theorem Proofs.Autograd.TapeNodes.ShapeOps.castVec_proof_irrel {n m : ℕ} (h₁ h₂ : n = m) (v : Vec n) : castVec h₁ v = castVec h₂ v

castVec is proof-irrelevant in its equality argument.

reshape is linear: on vectors it is just a type cast along Spec.Shape.size equality. We implement it as a Node to keep the DAG theorem applicable.

noncomputable def Proofs.Autograd.TapeNodes.ShapeOps.reshape {Γ : List Spec.Shape} {s₁ s₂ : Spec.Shape} (idx : Idx Γ s₁) (h : s₁.size = s₂.size) : Node Γ s₂

reshape node: reinterpret the same underlying coordinates as a different shape.

This is only definable when Spec.Shape.size s₁ = Spec.Shape.size s₂; at the vector level it is a cast.

PyTorch analogue: view/reshape operations that do not change the total number of elements. https://pytorch.org/docs/stable/tensor_view.html

noncomputable def Proofs.Autograd.TapeNodes.ShapeOps.reshapeFderiv {Γ : List Spec.Shape} {s₁ s₂ : Spec.Shape} (idx : Idx Γ s₁) (h : s₁.size = s₂.size) : NodeFDerivCorrect (reshape idx h)

NodeFDerivCorrect for reshape (it is linear/isometric).

flatten is a specialization of reshape to the canonical vector shape (.dim (Spec.Shape.size s) .scalar).

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

flatten node: specialization of reshape to the canonical vector shape (.dim (Spec.Shape.size s) .scalar).

PyTorch analogue: flatten when applied to a contiguous tensor. https://pytorch.org/docs/stable/generated/torch.flatten.html

noncomputable def Proofs.Autograd.TapeNodes.ShapeOps.flattenFderiv {Γ : List Spec.Shape} {s : Spec.Shape} (idx : Idx Γ s) : NodeFDerivCorrect (flatten idx)

NodeFDerivCorrect for flatten.

noncomputable def Proofs.Autograd.TapeNodes.ShapeOps.reindexVec {n m : ℕ} (e : Fin n ≃ Fin m) : Vec n → Vec m

Reindex a vector along a Fin equivalence (coordinate permutation/renaming).

noncomputable def Proofs.Autograd.TapeNodes.ShapeOps.reindexLin {n m : ℕ} (e : Fin n ≃ Fin m) : Vec n →L[ℝ] Vec m

The linear map induced by reindexVec.

@[simp]

theorem Proofs.Autograd.TapeNodes.ShapeOps.reindexLin_apply {n m : ℕ} (e : Fin n ≃ Fin m) (v : Vec n) : (reindexLin e) v = reindexVec e v

theorem Proofs.Autograd.TapeNodes.ShapeOps.inner_reindex_left {n m : ℕ} (e : Fin n ≃ Fin m) (x : Vec n) (y : Vec m) : inner ℝ (reindexVec e x) y = inner ℝ x (reindexVec e.symm y)

Move reindexVec across the left argument of an inner product.

def Proofs.Autograd.TapeNodes.ShapeOps.swapAdjacentEquiv (outer inner tail : ℕ) : Fin (outer * (inner * tail)) ≃ Fin (inner * (outer * tail))

Coordinate equivalence induced by exchanging two adjacent blocks.

def Proofs.Autograd.TapeNodes.ShapeOps.mapOuterEquiv (leading : ℕ) {source target : ℕ} (e : Fin source ≃ Fin target) : Fin (leading * source) ≃ Fin (leading * target)

Lift a coordinate equivalence pointwise under a leading axis.

noncomputable def Proofs.Autograd.TapeNodes.ShapeOps.reindex {Γ : List Spec.Shape} {source target : Spec.Shape} (idx : Idx Γ source) (e : Fin source.size ≃ Fin target.size) : Node Γ target

Reindex a tensor node by any coordinate equivalence. This is the proof-level primitive behind arbitrary axis permutations; its JVP uses the same permutation and its VJP uses the inverse.

noncomputable def Proofs.Autograd.TapeNodes.ShapeOps.reindexFDeriv {Γ : List Spec.Shape} {source target : Spec.Shape} (idx : Idx Γ source) (e : Fin source.size ≃ Fin target.size) : NodeFDerivCorrect (reindex idx e)

NodeFDerivCorrect for an arbitrary coordinate reindexing.