Softmax at arbitrary rank

Each output coordinate belongs to one last-axis row. Selecting that row is a linear map, so the vector softmax derivative applies after row selection at every tensor rank. The coordinate identities below use the actual max-shifted specification; no uniqueness assumption on the maximum is needed.

Coordinate functions give the finite-dimensional space on which we differentiate. Constructing the input tensor and reading the output are explicit in the derivative statement.

@[reducible, inline]

abbrev Proofs.Autograd.SoftmaxAxis.rowWidth : Spec.Shape → ℕ

Number of entries in a last-axis row; the scalar convention has one entry.

def Proofs.Autograd.SoftmaxAxis.lastCoordinate {s : Spec.Shape} : s.Coord → Fin (rowWidth s)

The last coordinate of a tensor coordinate, with index zero for a scalar.

def Proofs.Autograd.SoftmaxAxis.rowCoordinate {s : Spec.Shape} : s.Coord → Fin (rowWidth s) → s.Coord

Replace the last coordinate while keeping all leading coordinates fixed.

noncomputable def Proofs.Autograd.SoftmaxAxis.rowVector {s : Spec.Shape} (x : s.Coord → ℝ) (p : s.Coord) : Vec (rowWidth s)

Read one last-axis row from a tensor's coordinate function.

noncomputable def Proofs.Autograd.SoftmaxAxis.rowCLM {s : Spec.Shape} (p : s.Coord) : (s.Coord → ℝ) →L[ℝ] Vec (rowWidth s)

Row selection is linear and continuous on the finite coordinate space.

@[simp]

theorem Proofs.Autograd.SoftmaxAxis.rowCLM_apply {s : Spec.Shape} (p : s.Coord) (x : s.Coord → ℝ) : (rowCLM p) x = rowVector x p

Applying the row-selection map reads exactly that row.

theorem Proofs.Autograd.SoftmaxAxis.softmaxVec_one (x : Vec 1) (j : Fin 1) : (softmaxVec x).ofLp j = 1

The only coordinate of a one-entry softmax is one.

theorem Proofs.Autograd.SoftmaxAxis.rowVector_vector {n : ℕ} (x : TorchLean.Tensor ℝ [n]) (p : Spec.Shape.Coord [n]) : rowVector (fun (q : Spec.Shape.Coord [n]) => TorchLean.Tensor.Internal.Rep.get x q) p = getScalarE x

A rank-one tensor's selected row is its ordinary Euclidean vectorization.

theorem Proofs.Autograd.SoftmaxAxis.rowVector_unstack {n m : ℕ} {s : Spec.Shape} (x : TorchLean.Tensor ℝ (Spec.Shape.dim n (Spec.Shape.dim m s))) (p : (Spec.Shape.dim n (Spec.Shape.dim m s)).Coord) : rowVector (fun (q : (Spec.Shape.dim n (Spec.Shape.dim m s)).Coord) => TorchLean.Tensor.Internal.Rep.get x q) p = rowVector (fun (q : (Spec.Shape.dim m s).Coord) => TorchLean.Tensor.Internal.Rep.get (x.unstack p.1) q) p.2

Selecting a row after taking a leading slice selects the same coordinates.

theorem Proofs.Autograd.SoftmaxAxis.softmaxInnermostSpec_apply {s : Spec.Shape} (x : TorchLean.Tensor ℝ s) (p : s.Coord) : TorchLean.Tensor.Internal.Rep.get (Activation.Internal.softmaxInnermostSpec x) p = (softmaxVec (rowVector (fun (q : s.Coord) => TorchLean.Tensor.Internal.Rep.get x q) p)).ofLp (lastCoordinate p)

The stable innermost specification is the analytic vector softmax on the selected row.

The proof reduces each leading axis structurally, then uses the proved cancellation of the max shift in the vector kernel. Empty axes require no special numerical convention: they have no output coordinate at which the equality could fail.

noncomputable def Proofs.Autograd.SoftmaxAxis.innermostDeriv {s : Spec.Shape} (x : s.Coord → ℝ) : (s.Coord → ℝ) →L[ℝ] s.Coord → ℝ

Derivative of the innermost specification, assembled from its row derivatives.

theorem Proofs.Autograd.SoftmaxAxis.hasFDerivAt_softmaxInnermostSpec {s : Spec.Shape} (x : s.Coord → ℝ) : HasFDerivAt (fun (z : s.Coord → ℝ) (p : TorchLean.Tensor.Internal.Coord s) => TorchLean.Tensor.Internal.Rep.get (Activation.Internal.softmaxInnermostSpec (TorchLean.Tensor.Internal.Rep.ofFn z)) p) (innermostDeriv x) x

The max-shifted softmax specification is Fréchet differentiable at every tensor rank.

theorem Proofs.Autograd.SoftmaxAxis.softmaxJvp_one (x dx : Vec 1) (j : Fin 1) : (softmaxJvp x dx).ofLp j = 0

The one-entry softmax has zero differential.

theorem Proofs.Autograd.SoftmaxAxis.softmaxInnermostBackwardSpec_apply {s : Spec.Shape} (x dx : TorchLean.Tensor ℝ s) (p : s.Coord) : TorchLean.Tensor.Internal.Rep.get (Activation.Internal.softmaxInnermostBackwardSpec x dx) p = (softmaxJvp (rowVector (fun (q : s.Coord) => TorchLean.Tensor.Internal.Rep.get x q) p) (rowVector (fun (q : s.Coord) => TorchLean.Tensor.Internal.Rep.get dx q) p)).ofLp (lastCoordinate p)

The supplied innermost backward rule is the vector JVP on each selected row.

theorem Proofs.Autograd.SoftmaxAxis.innermostDeriv_apply {s : Spec.Shape} (x dx : s.Coord → ℝ) : (innermostDeriv x) dx = fun (p : TorchLean.Tensor.Internal.Coord s) => TorchLean.Tensor.Internal.Rep.get (Activation.Internal.softmaxInnermostBackwardSpec (TorchLean.Tensor.Internal.Rep.ofFn x) (TorchLean.Tensor.Internal.Rep.ofFn dx)) p

Applying the analytic derivative computes the actual specification's backward formula.

theorem Proofs.Autograd.SoftmaxAxis.fderiv_softmaxInnermostSpec {s : Spec.Shape} (x dx : s.Coord → ℝ) : (fderiv ℝ (fun (z : s.Coord → ℝ) (p : TorchLean.Tensor.Internal.Coord s) => TorchLean.Tensor.Internal.Rep.get (Activation.Internal.softmaxInnermostSpec (TorchLean.Tensor.Internal.Rep.ofFn z)) p) x) dx = fun (p : TorchLean.Tensor.Internal.Coord s) => TorchLean.Tensor.Internal.Rep.get (Activation.Internal.softmaxInnermostBackwardSpec (TorchLean.Tensor.Internal.Rep.ofFn x) (TorchLean.Tensor.Internal.Rep.ofFn dx)) p

The actual last-axis softmax derivative is its supplied backward rule at every rank.

theorem Proofs.Autograd.SoftmaxAxis.dot_softmaxInnermostBackwardSpec {s : Spec.Shape} (x dx δ : TorchLean.Tensor ℝ s) : Spec.dot (Activation.Internal.softmaxInnermostBackwardSpec x dx) δ = Spec.dot dx (Activation.Internal.softmaxInnermostBackwardSpec x δ)

The rank-general backward map is self-adjoint under the tensor inner product.