TensorOps
Package: sk.ainet.lang.tensor.ops
Modality: Core
matmul
Parameters
-
a: Tensorleft operand, shape[…, m, k]. The last dimensionkmust match the second-to-last dimension of [b]. -
b: Tensorright operand, shape[…, k, n]. Leading dimensions are broadcast against [a] using the usual broadcasting rules.
Definition
Given two matrices \(A \in \mathbb{R}^{m \times k}\) and \(B \in \mathbb{R}^{k \times n}\), the matrix product \(C = AB\) is defined as:
Where \(C \in \mathbb{R}^{m \times n}\), \(i\) ranges over rows \(1..m\), \(j\) over columns \(1..n\), and \(l\) is the summation index over the shared dimension \(k\).
Intuition
Matrix multiplication composes two linear transformations: each output element is the dot product of a row of \(A\) with a column of \(B\). It is the core primitive behind fully-connected layers, attention projections, and any linear map in a neural network’s forward pass.
Key properties:
-
Associativity: \((AB)C = A(BC)\)
-
Distributivity: \(A(B + C) = AB + AC\)
-
Non-commutativity: in general \(AB \neq BA\)
-
Identity: \(AI = IA = A\)
Complexity:
-
Standard algorithm: \(O(mnk)\)
-
Strassen’s algorithm: \(O(n^{2.807})\) for square matrices
-
Current theoretical best: \(O(n^{2.373})\)
Examples
val a: Tensor<FP32, FloatArray> = tensor(shape(2, 3)) { ... }
val b: Tensor<FP32, FloatArray> = tensor(shape(3, 4)) { ... }
val c = ops.matmul(a, b) // shape(2, 4)scaledDotProductAttention
Signature
fun scaledDotProductAttention(query:Tensor, key:Tensor, value:Tensor, mask:Tensor, scale:Float, causal:Boolean): TensorParameters
-
query: Tensor[batch, nHeads, seqLen, headDim] -
key: Tensor[batch, nKVHeads, kvLen, headDim] -
value: Tensor[batch, nKVHeads, kvLen, headDim] -
mask: Tensoroptional additive mask [batch, 1, seqLen, kvLen] (e.g. causal) -
scale: Floatscaling factor, defaults to 1/sqrt(headDim) -
causal: Booleanif true, apply causal masking (ignore [mask] parameter)