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.

theorem Proofs.TensorAxis.swapAdjacentAxes_apply {α : Type} [TorchLean.Storage α] {s : Spec.Shape} (x : TorchLean.Tensor α s) (depth : ℕ) (p : (s.swapAdjacentAtDepth depth).Coord) : TorchLean.Tensor.Internal.Rep.get (x.swapAdjacentAxes depth) p = TorchLean.Tensor.Internal.Rep.get x (swapCoordinate s depth p)

The tensor swap reads exactly the coordinate selected by swapCoordinate.

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.

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.

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.

def Proofs.TensorAxis.castCLM {s t : Spec.Shape} (h : s = t) : (s.Coord → ℝ) →L[ℝ] t.Coord → ℝ

Transporting coordinate functions through a shape equality is continuous and linear.

theorem Proofs.TensorAxis.castCLM_apply {s t : Spec.Shape} (h : s = t) (x : TorchLean.Tensor ℝ s) : ((castCLM h) fun (p : s.Coord) => TorchLean.Tensor.Internal.Rep.get x p) = fun (p : t.Coord) => TorchLean.Tensor.Internal.Rep.get (h ▸ x) p

The coordinate transport agrees with the shape transport used by the tensor API.

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.