Linear Algebra Helpers

Rank-polymorphic axis permutations, broadcasted matmul, and shape matching.

def TorchLean.Tensor.swapAdjacentAxes {β : Type} [Storage β] {shape : Spec.Shape} (tensor : Tensor β shape) (depth : ℕ) : Tensor β (shape.swapAdjacentAtDepth depth)

Swap adjacent tensor axes at depth and depth + 1.

def TorchLean.Tensor.Internal.swapAdjacentAxesCoordinate (shape : Spec.Shape) (depth : ℕ) : (shape.swapAdjacentAtDepth depth).Coord → shape.Coord

Map an output coordinate back to the input coordinate of an adjacent-axis swap.

At depth zero, an output coordinate (j, i, rest) reads (i, j, rest). At greater depths, the preceding coordinates stay fixed. If there are fewer than two axes left to exchange, the coordinate stays unchanged, as in Shape.swapAdjacentAtDepth.

def TorchLean.Tensor.Internal.swapAdjacentAxesDirect {β : Type} [Storage β] {shape : Spec.Shape} (tensor : Tensor β shape) (depth : ℕ) : Tensor β (shape.swapAdjacentAtDepth depth)

Build an adjacent-axis swap by reading each output entry directly from the input.

The recursive specification constructs tensor slices with unstack. Each slice owns a buffer, so constructing a column of a matrix this way repeatedly copies input rows. A coordinate pull constructs the output buffer in one pass and reads the corresponding scalar from the original tensor, without intermediate slices.

theorem TorchLean.Tensor.swapAdjacentAxes_apply {β : Type} [Storage β] {shape : Spec.Shape} (tensor : Tensor β shape) (depth : ℕ) (coordinate : (shape.swapAdjacentAtDepth depth).Coord) : Internal.Rep.get (tensor.swapAdjacentAxes depth) coordinate = Internal.Rep.get tensor (Internal.swapAdjacentAxesCoordinate shape depth coordinate)

Adjacent-axis swaps read the input at the coordinate with the same axes exchanged.

@[csimp]

theorem TorchLean.Tensor.swapAdjacentAxes_eq_swapAdjacentAxesDirect : @swapAdjacentAxes = @Internal.swapAdjacentAxesDirect

Compile the recursive specification as a direct coordinate pull.

Equality is proved for every scalar storage instance, shape, and swap depth. The public definition and its reduction equations remain available to proofs; compiled calls use the equivalent implementation that avoids copying intermediate slices.

def TorchLean.Tensor.permuteByAdjacentSwaps {β : Type} [Storage β] {s : Spec.Shape} (tensor : Tensor β s) (depths : List ℕ) : Tensor β (s.applyAdjacentSwaps depths)

Apply adjacent-axis swaps while retaining the resulting shape in the return type.

theorem TorchLean.Tensor.swapAdjacentAxes_zero {β : Type} [Storage β] {m n : ℕ} {s : Spec.Shape} (tensor : Tensor β (Spec.Shape.dim m (Spec.Shape.dim n s))) : tensor.swapAdjacentAxes 0 = dim fun (j : Fin n) => dim fun (i : Fin m) => get (get tensor i) j

Swapping at depth zero exchanges the two leading axes.

theorem TorchLean.Tensor.LinearAlgebra.Internal.rank_concat (left right : Spec.Shape) : (left.concat right).rank = left.rank + right.rank

Concatenated shapes add their ranks.

theorem TorchLean.Tensor.LinearAlgebra.Internal.sameRank_concat_right (left right suffix : Spec.Shape) (same : left.SameRank right) : (left.concat suffix).SameRank (right.concat suffix)

Appending the same suffix preserves equal ranks.

theorem TorchLean.Tensor.LinearAlgebra.Internal.extendBroadcastSuffix {source target : Spec.Shape} (suffix : Spec.Shape) (broadcast : source.CanBroadcastTo target) : (source.concat suffix).CanBroadcastTo (target.concat suffix)

Extend prefix-broadcast evidence across a fixed non-broadcasted tensor suffix.

theorem TorchLean.Tensor.LinearAlgebra.Internal.flattenBatchMatrix_size (batch : Spec.Shape) (m n : ℕ) : (batch.concat [m, n]).size = Spec.Shape.size [batch.size, m, n]

A batch prefix followed by a matrix has the size of the flattened batch matrix.

def TorchLean.Tensor.LinearAlgebra.Internal.matmulCommonBatchSpec {α : Type} [Storage α] [Add α] [Mul α] [Zero α] {batch : Spec.Shape} {m n p : ℕ} (A : Tensor α (batch.concat [m, n])) (B : Tensor α (batch.concat [n, p])) : Tensor α (batch.concat [m, p])

Multiply matrices that already share a common batch prefix.

@[simp]

theorem TorchLean.Tensor.LinearAlgebra.Internal.matmulCommonBatchSpec_scalar {α : Type} [Storage α] [Add α] [Mul α] [Zero α] {m n p : ℕ} (left : Tensor α [m, n]) (right : Tensor α [n, p]) : matmulCommonBatchSpec left right = matMulSpec left right

With no batch axes left, batched matmul is plain matmul.

@[simp]

theorem TorchLean.Tensor.LinearAlgebra.Internal.matmulCommonBatchSpec_dim {α : Type} [Storage α] [Add α] [Mul α] [Zero α] {count m n p : ℕ} {rest : Spec.Shape} (left : Fin count → Tensor α (rest.concat [m, n])) (right : Fin count → Tensor α (rest.concat [n, p])) : matmulCommonBatchSpec (dim left) (dim right) = dim fun (index : Fin count) => matmulCommonBatchSpec (left index) (right index)

Batched matmul over an outer axis is the batched matmul of each pair of slices.

def TorchLean.Tensor.matmulSpec {α : Type} [Storage α] [Add α] [Mul α] [Zero α] {batchA batchB batch : Spec.Shape} {m n p : ℕ} (broadcastA : batchA.CanBroadcastTo batch) (broadcastB : batchB.CanBroadcastTo batch) (A : Tensor α (batchA.concat [m, n])) (B : Tensor α (batchB.concat [n, p])) : Tensor α (batch.concat [m, p])

Matrix-rank matmul with explicit broadcasting of both batch prefixes.

A has shape batchA ++ [m, n], B has shape batchB ++ [n, p], and both batch prefixes broadcast to batch. The result has shape batch ++ [m, p].

def TorchLean.Tensor.matmulBackwardSpec {α : Type} [Storage α] [Add α] [Mul α] [Zero α] {batchA batchB batch : Spec.Shape} {m n p : ℕ} (broadcastA : batchA.CanBroadcastTo batch) (broadcastB : batchB.CanBroadcastTo batch) (A : Tensor α (batchA.concat [m, n])) (B : Tensor α (batchB.concat [n, p])) (dC : Tensor α (batch.concat [m, p])) : Tensor α (batchA.concat [m, n]) × Tensor α (batchB.concat [n, p])

Reverse-mode derivatives for matrix-rank matmul with broadcasted batch prefixes.