Last-axis softmax and log-softmax tape nodes

Softmax and log-softmax over matrix rows, together with their Jacobian/VJP lemmas and NodeFDerivCorrect wrappers for graph-level autograd proofs.

@[reducible, inline]

abbrev Proofs.Autograd.TapeNodes.SoftmaxLastAxis.MNSize (m n : ℕ) : ℕ

Flattened size for an m×n matrix when viewed as a single vector (m*n).

noncomputable def Proofs.Autograd.TapeNodes.SoftmaxLastAxis.rows {m n : ℕ} (x : Vec (MNSize m n)) : Fin m → Vec n

Split a flattened m*n vector into m rows of length n.

noncomputable def Proofs.Autograd.TapeNodes.SoftmaxLastAxis.unrows {m n : ℕ} (r : Fin m → Vec n) : Vec (MNSize m n)

Inverse of rows: assemble m rows back into a flattened m*n vector.

@[simp]

theorem Proofs.Autograd.TapeNodes.SoftmaxLastAxis.divNat_finProdFinEquiv {m n : ℕ} (p : Fin m × Fin n) : (finProdFinEquiv p).divNat = p.1

Row recovered from a flat index built by finProdFinEquiv.

Together with modNat_finProdFinEquiv this is the row-major convention spelled out: pairing then dividing gets you back where you started. Both are simp lemmas so rows/unrows round-trip without any manual index juggling.

@[simp]

theorem Proofs.Autograd.TapeNodes.SoftmaxLastAxis.modNat_finProdFinEquiv {m n : ℕ} (p : Fin m × Fin n) : (finProdFinEquiv p).modNat = p.2

Column recovered from a flat index built by finProdFinEquiv.

noncomputable def Proofs.Autograd.TapeNodes.SoftmaxLastAxis.rowsCLM {m n : ℕ} : Vec (MNSize m n) →L[ℝ] Fin m → Vec n

Continuous-linear-map version of rows.

noncomputable def Proofs.Autograd.TapeNodes.SoftmaxLastAxis.unrowsCLM {m n : ℕ} : (Fin m → Vec n) →L[ℝ] Vec (MNSize m n)

Continuous-linear-map version of unrows.

noncomputable def Proofs.Autograd.TapeNodes.SoftmaxLastAxis.forwardMN {m n : ℕ} (x : Vec (MNSize m n)) : Vec (MNSize m n)

Apply softmaxVec independently to each row of an m×n matrix (flattened representation).

noncomputable def Proofs.Autograd.TapeNodes.SoftmaxLastAxis.jvpMN {m n : ℕ} (x dx : Vec (MNSize m n)) : Vec (MNSize m n)

JVP of forwardMN, computed rowwise.

noncomputable def Proofs.Autograd.TapeNodes.SoftmaxLastAxis.derivMN {m n : ℕ} (x : Vec (MNSize m n)) : Vec (MNSize m n) →L[ℝ] Vec (MNSize m n)

Derivative of forwardMN as a continuous linear map.

theorem Proofs.Autograd.TapeNodes.SoftmaxLastAxis.hasFDerivAt_forwardMN {m n : ℕ} (x : Vec (MNSize m n)) : HasFDerivAt forwardMN (derivMN x) x

HasFDerivAt statement for forwardMN, using the rowwise softmax derivative.

theorem Proofs.Autograd.TapeNodes.SoftmaxLastAxis.jvpMN_eq_derivMN {m n : ℕ} (x dx : Vec (MNSize m n)) : jvpMN x dx = (derivMN x) dx

JVP computed by jvpMN agrees with applying the derivative derivMN.

theorem Proofs.Autograd.TapeNodes.SoftmaxLastAxis.inner_jvpMN_comm {m n : ℕ} (x dx δ : Vec (MNSize m n)) : inner ℝ (jvpMN x dx) δ = inner ℝ dx (jvpMN x δ)

Symmetry property of the JVP under the inner product (rowwise softmaxJvp commutation).

@[reducible, inline]

abbrev Proofs.Autograd.TapeNodes.LogSoftmaxLastAxis.MNSize (m n : ℕ) : ℕ

Reuse the m*n flattened size from SoftmaxLastAxis.

noncomputable def Proofs.Autograd.TapeNodes.LogSoftmaxLastAxis.forwardMN {m n : ℕ} (x : Vec (MNSize m n)) : Vec (MNSize m n)

Apply logSoftmaxVec independently to each row (flattened representation).

noncomputable def Proofs.Autograd.TapeNodes.LogSoftmaxLastAxis.jvpMN {m n : ℕ} (x dx : Vec (MNSize m n)) : Vec (MNSize m n)

JVP of forwardMN, computed rowwise via logSoftmaxJvp.

noncomputable def Proofs.Autograd.TapeNodes.LogSoftmaxLastAxis.vjpMN {m n : ℕ} (x δ : Vec (MNSize m n)) : Vec (MNSize m n)

VJP of forwardMN, computed rowwise via logSoftmaxVjp.

noncomputable def Proofs.Autograd.TapeNodes.LogSoftmaxLastAxis.derivMN {m n : ℕ} (x : Vec (MNSize m n)) : Vec (MNSize m n) →L[ℝ] Vec (MNSize m n)

Derivative of forwardMN as a continuous linear map.

theorem Proofs.Autograd.TapeNodes.LogSoftmaxLastAxis.hasFDerivAt_forwardMN {m n : ℕ} (x : Vec (MNSize m n)) : HasFDerivAt forwardMN (derivMN x) x

HasFDerivAt statement for logSoftmaxVec applied rowwise.

theorem Proofs.Autograd.TapeNodes.LogSoftmaxLastAxis.jvpMN_eq_derivMN {m n : ℕ} (x dx : Vec (MNSize m n)) : jvpMN x dx = (derivMN x) dx

JVP computed by jvpMN agrees with applying the derivative derivMN.

theorem Proofs.Autograd.TapeNodes.LogSoftmaxLastAxis.inner_jvpMN_vjp {m n : ℕ} (x dx δ : Vec (MNSize m n)) : inner ℝ (jvpMN x dx) δ = inner ℝ dx (vjpMN x δ)

logSoftmaxJvp / logSoftmaxVjp adjointness under the inner product, lifted rowwise.

noncomputable def Proofs.Autograd.TapeNodes.softmaxLast {Γ : List Spec.Shape} {m n : ℕ} (idx : Idx Γ (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) : Node Γ (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))

Tape node for applying softmaxVec along the last axis of an m×n matrix (rowwise).

noncomputable def Proofs.Autograd.TapeNodes.logSoftmaxLast {Γ : List Spec.Shape} {m n : ℕ} (idx : Idx Γ (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) : Node Γ (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))

Tape node for applying logSoftmaxVec along the last axis of an m×n matrix (rowwise).

noncomputable def Proofs.Autograd.TapeNodes.softmaxLastFderiv {Γ : List Spec.Shape} {m n : ℕ} (idx : Idx Γ (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) : NodeFDerivCorrect (softmaxLast idx)

NodeFDerivCorrect for softmaxLast (rowwise softmax).

noncomputable def Proofs.Autograd.TapeNodes.logSoftmaxLastFderiv {Γ : List Spec.Shape} {m n : ℕ} (idx : Idx Γ (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) : NodeFDerivCorrect (logSoftmaxLast idx)

NodeFDerivCorrect for logSoftmaxLast (rowwise log-softmax).