DAG Linear Algebra Primitives

Typed matrix products, vector-matrix products, and full sums for general computation graphs.

def NN.GraphSpec.DAG.PrimOp.matmul (batchA batchB batch : Spec.Shape) (mDim nDim pDim : ℕ) (broadcastA : batchA.CanBroadcastTo batch) (broadcastB : batchB.CanBroadcastTo batch) : PrimOp [batchA.concat [mDim, nDim], batchB.concat [nDim, pDim]] (batch.concat [mDim, pDim])

Multiply matrices after broadcasting their batch prefixes to a common shape.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.matmul_specFwd {batchA batchB batch : Spec.Shape} {mDim nDim pDim : ℕ} (broadcastA : batchA.CanBroadcastTo batch) (broadcastB : batchB.CanBroadcastTo batch) {α : Type} [TorchLean.Storage α] [Context α] (left : TorchLean.Tensor α (batchA.concat [mDim, nDim])) (right : TorchLean.Tensor α (batchB.concat [nDim, pDim])) : (matmul batchA batchB batch mDim nDim pDim broadcastA broadcastB).specFwd (TorchLean.TensorPack.cons left (TorchLean.TensorPack.cons right TorchLean.TensorPack.nil)) = TorchLean.Tensor.batchMatmulSpec broadcastA broadcastB left right

Pure evaluation of matrix multiplication with explicitly broadcast batch prefixes.

def NN.GraphSpec.DAG.PrimOp.Internal.vecMatCommonBatchSpec {α : Type} [TorchLean.Storage α] [Context α] {rows columns : ℕ} (batch : Spec.Shape) (vectors : TorchLean.Tensor α (batch.concat [rows])) (matrices : TorchLean.Tensor α (batch.concat [rows, columns])) : TorchLean.Tensor α (batch.concat [columns])

Multiply vectors and matrices pointwise over a common batch shape.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.Internal.vecMatCommonBatchSpec_scalar {α : Type} [TorchLean.Storage α] [Context α] {rows columns : ℕ} (vector : TorchLean.Tensor α [rows]) (matrix : TorchLean.Tensor α [rows, columns]) : vecMatCommonBatchSpec Spec.Shape.scalar vector matrix = Spec.vecMatMulSpec vector matrix

With no batch axes left, batched vector-matrix product is the plain one.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.Internal.vecMatCommonBatchSpec_dim {α : Type} [TorchLean.Storage α] [Context α] {count rows columns : ℕ} {rest : Spec.Shape} (vectors : Fin count → TorchLean.Tensor α (rest.concat [rows])) (matrices : Fin count → TorchLean.Tensor α (rest.concat [rows, columns])) : vecMatCommonBatchSpec (Spec.Shape.dim count rest) (TorchLean.Tensor.dim vectors) (TorchLean.Tensor.dim matrices) = TorchLean.Tensor.dim fun (index : Fin count) => vecMatCommonBatchSpec rest (vectors index) (matrices index)

A batch axis distributes over the vector-matrix product slicewise.

def NN.GraphSpec.DAG.PrimOp.broadcastVecMat (vectorBatch matrixBatch batch : Spec.Shape) (rows columns : ℕ) (broadcastVector : vectorBatch.CanBroadcastTo batch) (broadcastMatrix : matrixBatch.CanBroadcastTo batch) : PrimOp [vectorBatch.concat [rows], matrixBatch.concat [rows, columns]] (batch.concat [columns])

Multiply vectors by matrices after broadcasting their batch prefixes to a common shape.

@[simp]

theorem NN.GraphSpec.DAG.PrimOp.broadcastVecMat_specFwd {vectorBatch matrixBatch batch : Spec.Shape} {rows columns : ℕ} (broadcastVector : vectorBatch.CanBroadcastTo batch) (broadcastMatrix : matrixBatch.CanBroadcastTo batch) {α : Type} [TorchLean.Storage α] [Context α] (vector : TorchLean.Tensor α (vectorBatch.concat [rows])) (matrix : TorchLean.Tensor α (matrixBatch.concat [rows, columns])) : (broadcastVecMat vectorBatch matrixBatch batch rows columns broadcastVector broadcastMatrix).specFwd (TorchLean.TensorPack.cons vector (TorchLean.TensorPack.cons matrix TorchLean.TensorPack.nil)) = Internal.vecMatCommonBatchSpec batch (TorchLean.Tensor.broadcastTo ⋯ vector) (TorchLean.Tensor.broadcastTo ⋯ matrix)

Pure evaluation of broadcasted vector–matrix multiplication.

def NN.GraphSpec.DAG.PrimOp.sum (s : Spec.Shape) : PrimOp [s] Spec.Shape.scalar

Sum every scalar entry of a tensor.