TensorOps
Package: sk.ainet.lang.tensor.ops
Modality: Core
matmul
✅ DARC-validated by SKaiNET docs maintainers on 2026-05-24
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)conv3d
convTranspose1d
✖ Generated facts only (no human prose)
Signature
fun convTranspose1d(input:Tensor, weight:Tensor, bias:Tensor, stride:Int, padding:Int, outputPadding:Int, dilation:Int, groups:Int): TensorargMax
âš Prose present but not DARC-validated
Definition
Given a tensor \(X\) and a reduction dimension \(d\) of size \(n\), argMax returns, for
each position of the remaining dimensions, the index \(k^\*\) of the maximum value along \(d\):
Ties resolve to the lowest index (the \min above — numpy/greedy semantics). Dimension \(d\)
is removed from the output (no keepdim), and the result holds integer indices.
Intuition
argMax turns scores into a decision: "which entry won?" It is the final step of greedy decoding
(pick the highest-probability next token from a […, vocab] logits tensor) and of classification
(pick the top class). Because it selects an index rather than blending values, it is
non-differentiable and carries no backward rule.
SKaiNET keeps argMax a single op (like scaledDotProductAttention) and lowers it at the
StableHLO stage rather than adding a dedicated primitive. ArgMaxOperationsConverter composes it
from ops the exporter already emits:
-
stablehlo.iotaalong \(d\) — the candidate indices, -
reduce-
maximumalong \(d\), thenbroadcast_in_dimback — the per-position max value, -
compare EQof the input against that max — a boolean mask of the maxima, -
selectmask ? index : \(n\) (an out-of-range sentinel), then reduce-minimum— the lowest index among the maxima.
Output dtype. Indices are i32 in the compiled StableHLO. The eager CPU path materializes them
as index-valued floats in the input dtype, so the eager result is a portable Tensor<T, V> — an
i32 payload inside a float tensor is unreadable on Kotlin/Native and Wasm.
Examples
// [batch, vocab] logits -> [batch] token ids (greedy decode)
val logits: Tensor<FP32, FloatArray> = tensor(shape(1, 4)) { float(0.1f, 3.2f, 0.5f, 2.0f) }
val ids = ops.argMax(logits, dim = -1) // shape(1); value 1 (the max, 3.2, is at index 1)References
-
numpy.argmax — the tie-breaking (lowest index) reference semantics.
-
StableHLO
reduceandiota— the primitives the lowering composes.
scaledDotProductAttention
✖ Generated facts only (no human prose)
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)