Linear axis permutations

The softmax axis wrapper moves the selected axis to the last position, applies the row kernel, and restores the axes. Each move only changes which input coordinate an output coordinate reads. Its derivative therefore applies the same permutation to the input direction.

The coordinate maps below follow swapAdjacentAxes and permuteByAdjacentSwaps exactly. They also retain the executable convention for an invalid swap depth: that swap leaves the tensor unchanged. No nonempty-axis assumption is needed.

@[reducible, inline]

abbrev Proofs.TensorAxis.swapCoordinate (s : Spec.Shape) (depth : ℕ) : (s.swapAdjacentAtDepth depth).Coord → s.Coord

Source coordinate read by one adjacent-axis swap.

The derivative proofs and the compiled tensor operation use the same coordinate map. This name keeps the coordinate description available in the proof API.

def Proofs.TensorAxis.permutedCoordinate (s : Spec.Shape) (swaps : List ℕ) : (s.applyAdjacentSwaps swaps).Coord → s.Coord

Source coordinate read after a sequence of adjacent-axis swaps.

theorem Proofs.TensorAxis.permuteByAdjacentSwaps_apply {α : Type} [TorchLean.Storage α] {s : Spec.Shape} (x : TorchLean.Tensor α s) (swaps : List ℕ) (p : (s.applyAdjacentSwaps swaps).Coord) : TorchLean.Tensor.Internal.Rep.get (x.permuteByAdjacentSwaps swaps) p = TorchLean.Tensor.Internal.Rep.get x (permutedCoordinate s swaps p)

The coordinate description follows the actual tensor permutation at every rank.

noncomputable def Proofs.TensorAxis.permuteCLM {s : Spec.Shape} (swaps : List ℕ) : (s.Coord → ℝ) →L[ℝ] (s.applyAdjacentSwaps swaps).Coord → ℝ

Axis permutation on finite coordinate functions, bundled as a continuous linear map.

This is the pullback x ↦ x ∘ permutedCoordinate s swaps, which is linear for free (LinearMap.funLeft) and continuous because the domain is finite-dimensional.

theorem Proofs.TensorAxis.permuteCLM_apply {s : Spec.Shape} (swaps : List ℕ) (x : s.Coord → ℝ) : (permuteCLM swaps) x = fun (p : TorchLean.Tensor.Internal.Coord (s.applyAdjacentSwaps swaps)) => TorchLean.Tensor.Internal.Rep.get (TorchLean.Tensor.permuteByAdjacentSwaps (TorchLean.Tensor.Internal.Rep.ofFn x) swaps) p

The continuous linear map computes the library's tensor permutation.

theorem Proofs.TensorAxis.hasFDerivAt_permuteByAdjacentSwaps {s : Spec.Shape} (swaps : List ℕ) (x : s.Coord → ℝ) : HasFDerivAt (fun (z : s.Coord → ℝ) (p : TorchLean.Tensor.Internal.Coord (s.applyAdjacentSwaps swaps)) => TorchLean.Tensor.Internal.Rep.get (TorchLean.Tensor.permuteByAdjacentSwaps (TorchLean.Tensor.Internal.Rep.ofFn z) swaps) p) (permuteCLM swaps) x

Differentiating the actual tensor permutation permutes its input direction.