tensai

README

tensai - a tiny machine-learning framework in Go

tensai is a small machine-learning framework for learning and experiments. It implements forward passes, backpropagation, and optimization in pure Go; the default build has no external dependencies (the optional wgpu build tag adds exactly one, cgo-free: ebitengine/purego).

Documentation: mattn.github.io/tensai — guides for every package, in English and 日本語.

Features

• Matrix operations - Matrix plus basic operations such as Dot, Add, T, and AddBias. Tensors are float32 (tensai.Float)
• N-dimensional tensors - Tensor generalizes Matrix to any rank: element-wise Add/Sub/Mul/Div with NumPy-style broadcasting, batched MatMul (the leading axes broadcast, the per-matrix products run on the same kernel as Dot, parallelized across the batch), axis-permuting Transpose, Reshape with -1 inference, and zero-copy views to and from Matrix
• SIMD acceleration - AVX2 kernels written with Go's experimental simd/archsimd package: still pure Go, no cgo, no assembly files. Matmul, ReLU/LeakyReLU, Sigmoid/Tanh/Softmax (via a vectorized polynomial exp), GELU (via a vectorized erf), LayerNorm, and the Adam update are all 8-lane vectorized. Build with GOEXPERIMENT=simd on amd64 (Go 1.26 and 1.27 APIs both supported via build tags); every other build uses the portable fallbacks automatically
• Low-allocation training - layers reuse their forward/backward scratch buffers across training steps (a full MLP step runs in ~29 allocations), so GC stays out of the training loop; Predict always returns freshly allocated results
• Layers - Embedding, Dense, Conv2D, MaxPool2D, BatchNorm, LayerNorm, Dropout, plus ReLU, LeakyReLU, GELU, Sigmoid, Tanh, and Softmax activations
• WebGPU backend (experimental) - build with -tags wgpu (linux, macOS, Windows) and gpu.Open() runs batched MatMul as a WGSL compute shader on any GPU wgpu-native reaches (Vulkan, Metal, D3D12 — AMD, Intel, Apple, NVIDIA). The bindings go through ebitengine/purego, so there is still no cgo and no C compiler: the wgpu-native shared library is dlopen-ed at runtime
• int8 / int4 quantization - quant.Quantize / quant.Quantize4 build weight-only quantized twins: int4 group-wise with float32 accumulation, and int8 as a full integer path — weights in interleaved row quads, activations dynamically quantized to 7 bits, and the whole dot product running on the 256-bit u8 x s8 pairwise multiply-add plus a widening pair-add — two instructions per column, four rows deep — which reaches memory bandwidth (~31GB/s of weights on 16 cores). int4 halves the weights again — the difference between a 7B model fitting in RAM or not
• Loss functions - MeanSquaredError for regression, SoftmaxCrossEntropy for multi-class classification, and BinaryCrossEntropy for binary targets
• Optimizers - momentum SGD, Adam, and AdamW (decoupled weight decay)
• k-NN baseline - a knn.Classifier whose distance matrix runs on the same SIMD matmul kernel; useful as a no-training baseline next to the networks
• Dataset utilities - Dataset pairs inputs with targets and provides Shuffle, train/test Split (copy-free views), buffer-reusing mini-batch iteration with Batches, and Standardize/StandardizeWith
• Sequential models - stack layers and run Compile -> Fit / FitStep -> Predict
• Automatic differentiation - a micrograd-style reverse-mode autograd engine over matrices (Param / Input / Backward), for models that don't fit the Sequential mold; ToDot renders the computation graph for Graphviz
• Recurrence and attention - rnn.Cell, rnn.LSTMCell, and single-head rnn.SelfAttention built on the autograd engine, with backpropagation through time handled automatically
• Serialization - Save/Load (and SaveFile/LoadFile) round-trip trained Sequential parameters as JSON, including BatchNorm running statistics; SaveParams/LoadParams do the same for autograd parameters (RNN/LSTM/attention cells)
• TFLite export - the encoding/tflite package marshals Sequential models (FP32, NHWC) into .tflite flatbuffers that run on the TFLite/LiteRT runtimes and go-tflite, with the FlatBuffers writer implemented in-tree — still no dependencies
• safetensors - encoding/safetensors reads the checkpoint format most published model weights ship in — lazily, one tensor at a time, with F16/BF16/F64 converted to float32 — and writes F32 checkpoints; interoperability is verified against the reference implementation in both directions. Also dependency-free
• GGUF - encoding/gguf reads llama.cpp's model container — typed metadata plus lazily-loaded tensors, with F16/BF16, the block-quantized Q8_0/Q4_0/Q4_1/Q5_0/Q5_1, the K-quants Q2_K through Q6_K, IQ4_NL, and gpt-oss's MXFP4 dequantized to float32 — verified block-exact against real llama.cpp conversions. Dependency-free as well
• ONNX export - encoding/onnx marshals Sequential models into ONNX (opset 13, FP32) with a hand-written protobuf encoder; onnxruntime reproduces Predict to ~1e-7 relative error. ONNX convolutions are NCHW, which is tensai's own row layout, so nothing is reordered
• Tokenizers - the tokenizer package loads Hugging Face tokenizer.json files and implements the byte-level BPE family (GPT-2, Llama 3, Qwen, ...), including the split patterns Go's regexp cannot express, as hand-written scanners — the GPT-2, cl100k, and o200k (gpt-4o/gpt-oss) families — plus SentencePiece (Gemma, the Llama-2 era) built from GGUF vocabularies via NewSPM; encodings match the reference tokenizers library and llama-tokenize exactly

Layout

go.mod              Module definition (github.com/mattn/tensai)
.                   Core: Float, Matrix, N-d Tensor, Dot and the AVX2/portable matmul kernels
layer               Layer interface, Dense, Conv2D, MaxPool2D, BatchNorm, LayerNorm, Dropout, Embedding, activations
loss                Loss functions (MSE, SoftmaxCrossEntropy, BCE)
optim               Optimizers (SGD, Adam, AdamW)
model               Sequential model, training loop, and JSON save/load
autograd            Reverse-mode automatic differentiation (Node graph), Trainer, parameter save/load
rnn                 rnn.Cell / LSTMCell / SelfAttention on the autograd engine
knn                 k-NN baseline classifier
dataset             Shuffle, split, standardize, mini-batch iteration
quant               int8 / int4 / grouped-int8 / MXFP4 weight-only quantization
gpu                 WebGPU backend via purego + wgpu-native (build tags wgpu / wgpu24)
internal/kernels    Element-wise kernels: scalar bodies plus the AVX2 versions incl. vectorized exp
internal/simd       Load/store shims over both simd/archsimd API generations
internal/dims       Shape arithmetic shared between the core and the GPU backend
_example/helloworld Smallest possible program: add two values on the graph
_example/dataset    Dataset workflow: shuffle, split, standardize, batches
_example/xor        Runnable XOR training example
_example/fizzbuzz   Runnable FizzBuzz classification example
_example/spiral     Runnable 3-class spiral classification example
_example/iris       Runnable Iris classification example
_example/mnist      Runnable MNIST classifier (-model dense, cnn, or knn) with save/load
_example/charrnn    Character-level LSTM text generation on the autograd engine
_example/plasma     Demoscene-style terminal plasma rendered by a neural network
_example/dot        Graphviz DOT export of the z = x + y graph
_example/tensor     Tour of the n-d Tensor: broadcasting, batched MatMul, attention
_example/wgpu       WebGPU MatMul: adapter info, CPU cross-check, GPU vs CPU sweep
_example/gpt2       The published GPT-2 (124M) checkpoint generating text in pure Go
_example/qwen       Qwen2.5-0.5B-Instruct chatting in pure Go: RoPE, GQA, SwiGLU
cmd/tensai          The tensai command: run, chat, and serve subcommands over internal/llm

Usage

Regression: learn XOR with MSE

net := model.NewSequential()
net.Add(layer.NewDense(8))
net.Add(&layer.Tanh{})
net.Add(layer.NewDense(1))
net.Add(&layer.Sigmoid{})

net.Compile(2, loss.MeanSquaredError{}, optim.NewAdam(0.05))
net.Fit(inputs, targets, 5000)

pred, _ := net.Predict(inputs)

Classification: softmax + cross-entropy

net := model.NewSequential()
net.Add(layer.NewDense(8))
net.Add(&layer.ReLU{})
net.Add(layer.NewDense(2)) // output width = number of classes

net.Compile(2, loss.SoftmaxCrossEntropy{}, optim.NewAdam(0.05))

SoftmaxCrossEntropy expects targets as an Mx1 matrix of class indices. Softmax is applied inside the loss, so Predict returns raw logits. Use argmax for classification.

Datasets

ds, _ := dataset.New(inputs, targets)
ds.Shuffle(rng)
train, test, _ := ds.Split(0.2)          // views, no copying
mean, std := train.Standardize()         // fit on train...
test.StandardizeWith(mean, std)          // ...apply to test

for epoch := 0; epoch < epochs; epoch++ {
	train.Batches(32, rng, func(in, tgt *tensai.Matrix) error {
		_, err := net.FitStep(in, tgt)
		return err
	})
}

Convolution, regularization, and saving

net := model.NewSequential()
net.Add(layer.NewConv2D(8, 3, 1, 1)) // outC, kernel, stride, pad
net.Add(&layer.ReLU{})
net.Add(layer.NewMaxPool2D(2))
net.Add(layer.NewDense(64))
net.Add(layer.NewBatchNorm())
net.Add(layer.NewLeakyReLU(0.01))
net.Add(layer.NewDropout(0.3))
net.Add(layer.NewDense(10))

// The input geometry is stated once; the spatial shape threads through
// the stack, so the conv and pool layers pick their dimensions up from it.
net.CompileImage(layer.Image{H: 28, W: 28, C: 1}, loss.SoftmaxCrossEntropy{}, optim.NewAdamW(0.001, 0.01))
net.Fit(inputs, targets, 10)

net.SaveFile("model.json")
// Later: build + Compile the same architecture, then
net.LoadFile("model.json")

Conv2D and MaxPool2D treat each row as a channel-major image: index = (channel*height + y)*width + x. Dropout and BatchNorm switch automatically between training behavior (inside Fit/FitStep) and inference behavior (inside Predict).

Embedding keeps the current matrix-only API: each input row is a token-id sequence, and the layer concatenates the looked-up embedding vectors across columns. For example, Compile(4, ...) plus NewEmbedding(vocab, 8) turns an Mx4 token-id matrix into an Mx32 dense feature matrix that can feed LayerNorm, GELU, and Dense.

Tokenizers

import "github.com/mattn/tensai/tokenizer"

tok, err := tokenizer.Load("tokenizer.json") // the file models ship on Hugging Face
ids := tok.Encode("Hello, I'm a language model,")
text := tok.Decode(ids)
eos, _ := tok.ID("<|endoftext|>")

Byte-level BPE as GPT-2, Llama 3, and Qwen use it. The pre-tokenization regexes these models declare need lookahead and inline case-insensitive groups that regexp cannot express, so the two patterns that exist in the wild — the GPT-2 split and the cl100k-style split — are hand-written scanners, and anything else is rejected rather than silently mis-tokenized. Special tokens are matched verbatim during encode. Verified against the reference tokenizers library: an adversarial corpus and 2000 fuzzed strings encode and decode identically for both GPT-2 and Qwen2.5 (see tokenizer/verify_hf.py). An NFC normalizer passes through — input is assumed already NFC, which virtually all real-world text is.

ONNX export

import tensaionnx "github.com/mattn/tensai/encoding/onnx"

err := tensaionnx.MarshalFile("model.onnx", model)

Same layer support as the TFLite export (Dense, Conv2D, MaxPool2D, BatchNorm folded to Mul+Add, Dropout dropped, Softmax on dense features, and the ReLU/LeakyReLU/Sigmoid/Tanh activations), but no layout gotcha: ONNX convolutions are NCHW, which is exactly tensai's channel-major row layout, so the exported model consumes the same flattened rows tensai does, as a [1, C, H, W] tensor. Verified against onnxruntime to ~1e-7 relative error (see encoding/onnx/verify_onnxruntime.py).

TFLite export

import tensaitflite "github.com/mattn/tensai/encoding/tflite"

// after training:
err := tensaitflite.MarshalFile("model.tflite", model)

Supported layers: Dense, Conv2D (VALID/SAME padding), MaxPool2D, BatchNorm (folded into Mul+Add), Dropout (dropped), Softmax, and the ReLU/LeakyReLU/Sigmoid/Tanh activations. Exported convolutions follow TFLite's NHWC layout — feed the exported model NHWC input; weight reordering is handled during export. Outputs have been verified to match Predict to ~1e-7 relative error on the LiteRT interpreter (see encoding/tflite/verify_litert.py). Alias the import when combining with go-tflite, which also names its package tflite.

safetensors checkpoints

encoding/safetensors opens the format published model weights usually ship in. Open parses only the header; each Tensor call reads just that tensor's bytes, so single tensors come out of multi-gigabyte checkpoints without loading the rest. F32 loads as-is and F16/BF16/F64 convert to tensai's float32:

import "github.com/mattn/tensai/encoding/safetensors"

f, err := safetensors.Open("model.safetensors")
defer f.Close()
w, err := f.Tensor("model.layers.0.attention.wq.weight") // *tensai.Tensor

encoding/gguf does the same for llama.cpp's GGUF container: Open parses the typed metadata (String/Int/Float/KV) and the tensor directory, and each Tensor call reads and dequantizes just that tensor — F32/F16/BF16 plus the block-quantized Q8_0, Q4_0/Q4_1, Q5_0/Q5_1, the K-quants Q2_K through Q6_K, and the nonlinear IQ4_NL, so the whole ladder of checkpoints usually published for llama.cpp — q2_k up through q8_0 — opens directly. Dimensions come back row-major like every other reader here. One caveat inherited from the format: llama.cpp's converter permutes attention q/k projection rows into its interleaved RoPE order, which consumers pairing GGUF weights with half-split RoPE must undo.

Names, Info, and Metadata inspect a checkpoint without loading it; Save/SaveFile write F32 checkpoints that the reference implementation reads back bit-for-bit.

_example/gpt2 puts the reader to work on a real model: it downloads the published GPT-2 small (124M) checkpoint from Hugging Face, loads the weights through this package, tokenizes with a from-scratch byte-level BPE, and decodes with a KV cache — every matvec running on the same Dot kernel as the rest of tensai, at ~30 tok/s with the AVX2 build:

$ GOEXPERIMENT=simd go run ./_example/gpt2 -n 20
Hello, I'm a language model, not a programming language. I'm a language model. ...

The greedy continuation matches GPT-2's well-known reference output token for token, which pins the whole pipeline — reader, tokenizer, and forward pass — in one check.

The prompt runs through the model as one batched pass; with -gpu (built with -tags wgpu or wgpu24) every block's causal multi-head attention becomes a single masked dispatch on the GPU. A 600-token prompt prefills 1.5x faster even through dozen inside WSL2; native drivers gain more.

-q8 quantizes the decode-path weights to int8 (weight-only, per-column scales) and doubles generation — 23 to 46 tok/s on the same machine — because decode streams the whole checkpoint per token and int8 pulls a quarter of the bytes. The text stays coherent but greedy decoding no longer reproduces the float32 reference tokens exactly; use the default float32 path for the reference check.

_example/qwen does the same for modern instruction-tuned models: RMSNorm, rotary position embeddings, grouped-query attention, and a SwiGLU MLP, loaded from safetensors (config.json drives the dimensions, sharded checkpoints come through their index.json) or from a single llama.cpp GGUF that carries config, tokenizer, and weights in one file — -gguf qwen2.5-0.5b-instruct-q8_0.gguf -q8 chats with nothing else on disk. One runtime speaks nine architectures, each contributing its own twist:

family	models	what it adds
qwen2	Qwen 1.5/2/2.5, Qwen2.5-Coder, the R1-Distill-Qwen line	attention biases
qwen3	Qwen3 dense	per-head QK-norm, explicit head_dim, -think
llama	Llama 2/3, SmolLM2, Mistral, R1-Distill-Llama	the block everyone forked
smollm3	SmolLM3-3B	RoPE skipped every fourth layer
gemma3	Gemma 3	sliding windows on 5/6 layers, sandwich norms, gelu-tanh gate, SentencePiece
phi3	Phi-3/3.5-mini	q/k/v and gate/up shipped pre-fused
qwen2moe / qwen3moe	Qwen1.5-MoE-A2.7B, Qwen3-30B-A3B	top-k routed experts, a shared expert on qwen2moe
gpt-oss	gpt-oss-20b	MXFP4 experts, attention sinks, YaRN rope, harmony channels

The DeepSeek-R1 distills need no family of their own — they are stock qwen2/llama blocks wearing DeepSeek's turn markers, which the loader spots in the embedded chat template and switches automatically, <think> reasoning included. Mixture-of-experts blocks route each token through its top-k experts, repacked per expert straight from the GGUF's 3D tensors: Qwen1.5-MoE-A2.7B (14B total, 2.7B active) answers at ~9 tok/s from a 20-second load, and gpt-oss-20b — its experts kept in their native MXFP4 blocks, expanded through a one-shuffle table-lookup kernel — reasons in its harmony analysis channel and answers on the same 15GB machine.

With -q8/-q4 each weight quantizes as it loads and its float32 copy dies immediately, so the full-precision model never has to fit in memory, and the layers load in parallel with the quantizer splitting columns across CPUs. Quantized GGUF checkpoints skip the float32 detour entirely: Q8_0, Q4_0, Q5_0, the whole Q4_K/Q5_K/Q6_K K-quant family, and MXFP4 repack straight from the memory-mapped file — nibbles copy raw where the grids line up, five- and six-bit spans renormalize with integer rounding under -q4, everything widens onto a finer int8 grid under -q8 — keeping llama.cpp's own quantization intact. A 1.5B Q4_K_M loads in about 3 seconds instead of 8 (-requant restores the float detour, trading the much slower load for about 10% more decode speed from its coarser symmetric tables); a 3B Q8_0 opens in 5 seconds instead of 32.

The first -gguf load also writes the repacked weights to a cache file next to the model (-nocache opts out), and every later load just memory-maps it: the 1.5B Q4_K_M reopens in ~0.3 seconds, a Mistral 7B in well under a second, and gpt-oss-20b in under two. Beyond the instant reopen, mapped weights are clean file-backed pages the kernel can drop and re-read at will — on a machine where the model barely fits, that replaces swap thrashing with ordinary page cache behavior, which is the same structural advantage llama.cpp gets from decoding straight out of its mmap'd file.

On a 15GB machine the ladder looks like: 0.5B at ~40 tok/s with -q8 and a 1.5B Q4_K_M at ~25 with -q4 (the tiled integer kernels, measured on native Windows), and Qwen2.5-7B-Instruct — 15GB of BF16 shards, int4-quantized on the fly during a two-minute load into ~6GB resident — answering correctly at 3.5 tok/s. Prompts feed through a batched prefill: QMatrix.MatMul streams the weights once per block of eight token rows instead of once per token, cutting the wait before the first generated token by around 6x.

-draft points at a smaller same-family model for speculative decoding (greedy only): the draft proposes a few tokens, one batched pass of the big model verifies them, and rejections roll the caches back, so the output is exactly what the big model alone would produce — Qwen2.5-7B with the 0.5B drafting goes from 1.2 to 1.6 tok/s; the draft only pays off when the target is much larger than it. Sampling (-temp above 0) restricts itself to the nucleus: -topp 0.9 keeps the smallest probability-sorted set of tokens holding 90% of the mass, so the long tail where repetition loops live never gets a lottery ticket. -chat turns it into a multi-turn conversation on stdin — the KV cache carries the whole dialogue, so each turn only processes its own tokens — and -serve :8080 exposes the same model as an OpenAI-compatible /v1/chat/completions endpoint (messages array, SSE streaming, usage counts), so any OpenAI client pointed at it chats with a pure-Go model:

$ GOEXPERIMENT=simd go run ./_example/qwen -q8 -prompt "What is the capital of France?"
The capital of France is Paris.
43 tokens in 1.3s (33.1 tok/s)

Automatic differentiation

When a model doesn't fit the Sequential mold (weight sharing, custom losses, exotic architectures), build the computation directly and let reverse-mode autodiff derive the gradients:

w1 := autograd.Param(tensai.RandomMatrix(2, 8, rng))
b1 := autograd.Param(tensai.NewMatrix(1, 8))
w2 := autograd.Param(tensai.RandomMatrix(8, 1, rng))
trainer := autograd.NewTrainer(optim.NewAdam(0.05), w1, b1, w2)

for step := 0; step < 2000; step++ {
	loss := autograd.Input(x).MatMul(w1).AddRow(b1).Tanh().MatMul(w2).Sigmoid().MSELoss(y)
	trainer.Step(loss) // backward + update + zero grads, returns the loss value
}

For manual control, the pieces are still public: loss.Backward(), p.Grad, and autograd.ZeroGrads(params...).

The graph a loss node holds can be visualized: loss.ToDot() returns Graphviz DOT (label leaves with .Named("w1")), so go run ./_example/dot | dot -Tsvg > graph.svg draws the network the same way Gorgonia's encoding/dot does.

Graphs are built dynamically per step (define-by-run) and are single-use. Available ops: MatMul, Add, Sub, MulElem, Scale, AddRow, T, Softmax, ReLU, Sigmoid, Tanh, Sum, Mean, MSELoss, and SoftmaxCELoss. Shape mismatches panic during graph construction. Every op's gradient is verified against finite differences in the test suite.

Recurrent networks and attention

rnn.Cell, rnn.LSTMCell, and rnn.SelfAttention are built on the autograd engine, so unrolling a sequence is a plain Go loop and backpropagation through time comes for free:

cell := rnn.NewLSTMCell(inSize, hidden, rng)
wOut := autograd.Param(tensai.RandomMatrix(hidden, numClasses, rng))
bOut := autograd.Param(tensai.NewMatrix(1, numClasses))
trainer := autograd.NewTrainer(optim.NewAdam(0.01), append(cell.Params(), wOut, bOut)...)

for step := 0; step < epochs; step++ {
	h, c := cell.InitState(batch)
	for _, x := range steps { // one (batch x inSize) matrix per time step
		h, c = cell.Step(autograd.Input(x), h, c)
	}
	logits := h.MatMul(wOut).AddRow(bOut)
	trainer.Step(logits.SoftmaxCELoss(labels))
}

rnn.SelfAttention operates on one (seqLen x inSize) sequence node: attn.Forward(x) computes softmax(Q*K^T/sqrt(d))*V with learned projections; the raw rnn.Attention(q, k, v) form is also exposed.

Autograd parameters are saved and restored positionally with autograd.SaveParamsFile("cell.json", cell.Params()...) / autograd.LoadParamsFile("cell.json", cell.Params()...) — build the same cell, then load.

N-d tensors: broadcasting and batched MatMul

Tensor generalizes Matrix to any rank. Element-wise ops broadcast NumPy-style, and MatMul multiplies whole stacks of matrices at once — the leading batch axes broadcast too, so a shared 2-D weight applies to every sequence in a batch in one call:

x := tensai.NewTensor(4, 6, 3)                    // (batch, position, channel)
mean, _ := tensai.NewTensorFromSlice([]float32{0.5, -1, 2}, 3)
centered, _ := x.Sub(mean)                        // (4,6,3) - (3)   -> (4,6,3)
h, _ := tensai.MatMul(centered, w)                // (4,6,3) @ (3,8) -> (4,6,8)

kt, _ := k.Transpose()                            // swap the last two axes
scores, _ := tensai.MatMul(q, kt)                 // (4,6,8) @ (4,8,6) -> (4,6,6)
scores.Scale(1 / float32(math.Sqrt(8)))
out, _ := tensai.MatMul(scores, v)                // attention for the whole batch

Tensors are contiguous and row-major; Reshape (with -1 inference) and the Matrix/Tensor conversions are zero-copy views, while Transpose accepts an arbitrary axis permutation and materializes the result. See _example/tensor for the runnable version.

GPU MatMul over WebGPU (experimental)

Building with -tags wgpu (linux/darwin/windows) enables a GPU backend for batched MatMul with the same shape and broadcasting semantics as the CPU version:

dev, err := gpu.Open() // fails cleanly when no GPU / library is present
if err != nil { /* fall back to tensai.MatMul */ }
defer dev.Close()
fmt.Println(dev.Name()) // e.g. "AMD Radeon 780M (integrated)"
out, err := dev.MatMul(a, b)

On machines with both an integrated and a discrete GPU, pass a preference: gpu.Open(gpu.LowPower) steers to the iGPU, gpu.HighPerformance to the dGPU (it is a hint — with a single adapter you always get that one).

Buffers can also stay resident on the GPU, so a weight rides the bus once instead of on every call and intermediates never leave the device:

gw, _ := dev.Upload(w)              // weight uploaded once
defer gw.Free()                     // GPU memory is not garbage collected
gx, _ := dev.Upload(x)
h, _ := gx.MatMul(gw)               // chain freely; nothing touches the host
out, _ := h.MatMul(gw2)
result, _ := out.Download()         // one readback at the end

dev.MatMul(a, b) is shorthand for Upload → MatMul → Download → Free. Residency matters most on discrete GPUs, where every transfer crosses PCIe; on shared-memory iGPUs the win is smaller and comes mainly from skipping intermediate readbacks.

Beyond MatMul, resident tensors support MatMulT (multiply by a transposed operand without materializing the transpose), an in-place Scale, and a row-parallel Softmax over the last axis — enough to run single-head attention entirely on the GPU:

out, _ := gq.Attention(gk, gv)                 // softmax(q@k^T/sqrt(d))@v, no host round-trips
out, _ = gq.MultiHeadAttention(gk, gv, heads)  // packed (batch, seq, heads*dh) layout

Multi-head attention carves each head out of the packed layout with strided kernels — the matmul kernels take explicit row strides and per-batch offsets — so no permute is ever materialized. The causal variants (CausalAttention, CausalMultiHeadAttention) mask future positions inside the kernel, with k and v allowed to hold more positions than q — the prompt-prefill and chunked-decode patterns of autoregressive models — so no mask tensor is ever built either. CausalMultiHeadAttention runs as one fused flash-attention-style dispatch (an online softmax over kv tiles, for head dimensions up to 128): the scores matrix never exists, so memory stays at q+k+v+output regardless of sequence length, and shapes whose scores would blow past the device's storage-buffer limit — batch 8, heads 8, seq 1024 is a 256MiB scores matrix — just run.

There is no cgo involved: the bindings load the wgpu-native shared library at runtime via ebitengine/purego (dlopen on linux/macOS, LoadLibrary on Windows). Download a v22.1.0.5 release binary (the C API these bindings target), then either install it where the loader finds it or point TENSAI_WGPU_LIB at it:

curl -sLO https://github.com/gfx-rs/wgpu-native/releases/download/v22.1.0.5/wgpu-linux-x86_64-release.zip
unzip wgpu-linux-x86_64-release.zip -d wgpu
TENSAI_WGPU_LIB=$PWD/wgpu/lib/libwgpu_native.so go test -tags wgpu ./...

On Windows, take wgpu-windows-x86_64-msvc-release.zip from the same release and point the variable at the wgpu_native.dll inside it (any wgpu_native.dll on PATH or next to the executable is found without the variable):

$env:TENSAI_WGPU_LIB="$PWD\wgpu\lib\wgpu_native.dll"
go run -tags wgpu ./_example/wgpu

_example/wgpu -sweep walks a ladder of sizes and marks where the GPU overtakes the CPU kernel. It reports gpu+xfer for the convenient Upload → MatMul → Download call and resident for inputs uploaded once and reused (the final result is still downloaded each iteration). Because the CPU side is the same dotRows kernel the rest of the package uses, building the example twice compares portable Go, AVX2, and both GPU usage patterns:

GOEXPERIMENT=nosimd go build -tags wgpu -o wgpu-nosimd ./_example/wgpu
GOEXPERIMENT=simd   go build -tags wgpu -o wgpu-simd   ./_example/wgpu
./wgpu-nosimd -sweep && ./wgpu-simd -sweep

The crossover moves with the CPU kernel, GPU driver, and transfer pattern. The res/cpu column and crossover marker use the resident-input timing, since that is the normal pattern for repeated inference. On a Ryzen iGPU (AMD Radeon 780M, native Windows, AVX2 CPU kernel) the register-tiled kernels put every rung of the ladder on the GPU side:

             shape                   MFLOP   gpu+xfer   resident        cpu   res/cpu
mnist dense  1x100x784@784x128        20.1     1.51ms      597µs      652µs     1.09x
mnist conv2  1x19600x72@72x16         45.2    1.388ms      763µs    2.354ms     3.09x
tiny         1x128x128@128x128         4.2      432µs      302µs      410µs     1.36x
small        1x512x512@512x512       268.4    1.331ms    1.216ms    6.865ms     5.65x
medium       8x512x512@512x512      2147.5    8.053ms    6.297ms    71.56ms    11.36x
large        32x512x512@512x512     8589.9    86.71ms   28.856ms  266.277ms     9.23x
huge         64x512x512@512x512    17179.9  116.374ms   62.128ms  566.726ms     9.12x

Arithmetic no longer dominates the convenient path — gpu+xfer at large spends two thirds of its time on the bus — which is exactly what keeping inputs resident is for. Through a translation layer like dozen inside WSL2 the ratios shrink to roughly parity-to-3x, and on CPU Vulkan implementations the GPU path loses outright; measure on the driver you will ship on.

Quantized weights stay quantized on the device too: UploadQ8 packs a QMatrix four int8 weights per u32, and gpu.QMatrix.MatMul dequantizes them in registers, so a decode matvec — whose cost is streaming the weights — moves a quarter of the f32 bytes. On the same iGPU through dozen it runs the matvec 2.2x faster than the resident f32 kernel.

UploadQ4 does the same for the int4 twin — nibbles packed four row-pair bytes per u32, group scales folded at group boundaries in registers — so -q4 -gpu runs models whose int8 weights would not fit. The rest of a transformer decode step is there as well — RMSNorm, in-place RoPE, Add, SiluMul, GroupedCausalAttention (a KV cache packing fewer heads than the queries, read up to a valid length), and CopyRowsInto to append fresh k/v rows to a resident cache — so _example/qwen -q8 -gpu runs every block on the device and only the hidden state comes back per token. BeginBatch/Flush record a whole token's dispatches into one submission, and freed intermediates recycle through a buffer pool, which together took a dozen-translated decode from 1.2 to ~17 tok/s steady state on the machine above. On native Windows the same iGPU speaks D3D12 directly, and -q8 -gpu held the 0.5B decode crown for a while — 29.7 tok/s against 23.2 on the AVX2 path — until the tiled integer kernels took it back: the CPU now decodes the same model at ~42 tok/s against ~30 on the GPU, which stays useful for keeping the cores free.

wgpu-native picks Vulkan on Linux, Vulkan or D3D12 on Windows, and Metal on macOS, so AMD, Intel, Apple, and NVIDIA GPUs all work — as do CPU Vulkan implementations like lavapipe, which is how the tests run on machines without a GPU. gpu.MatMul uploads the operands and reads the product back on every call; Upload plus gpu.Tensor.MatMul keeps inputs and intermediates resident, so only the final result needs to cross the bus. Without the build tag gpu.Open returns an error and nothing else changes.

-tags wgpu24: the new wgpu-native API, and the real GPU inside WSL2

-tags wgpu24 (linux/darwin/windows) builds the same gpu.Open API against the reworked wgpu-native C API instead — pair it with a v29-series release binary. The new API's payoff is WGPUInstanceFlag_AllowUnderlyingNoncompliantAdapter, which un-hides non-conformant Vulkan drivers. Concretely: Mesa's dozen (Vulkan-on-D3D12, shipped in the kisak-mesa PPA) exposes the real host GPU inside WSL2, but the v22 API hides it as non-conformant and falls back to lavapipe; the wgpu24 build reaches it:

VK_DRIVER_FILES=/path/to/dzn_icd.json \
TENSAI_WGPU_LIB=$PWD/wgpu29/lib/libwgpu_native.so \
    go run -tags wgpu24 ./_example/wgpu   # adapter: Microsoft Direct3D12 (AMD Radeon(TM) Graphics)

The new API passes structs by value. Every one of them is reached through a pointer field except the three callback-info arguments, and those are the only per-OS code in the binding: wgpu24_callinfo.go hands the 40-byte struct to SysV/AAPCS in registers, while wgpu24_callinfo_windows.go passes its address, because the Windows x64 convention already defines any aggregate that is not 1, 2, 4, or 8 bytes wide as passed by reference. WGPUFuture results come back in RAX either way. When both tags are set, wgpu24 wins.

On Windows, pair it with wgpu-windows-x86_64-msvc-release.zip from the same v29 release:

$env:TENSAI_WGPU_LIB="$PWD\wgpu29\lib\wgpu_native.dll"
go run -tags wgpu24 ./_example/wgpu

Note that new does not mean faster: on a Radeon 780M at 32x512x512@512x512 the v22 library runs the same shader in 85ms and the v29 one in 165ms (D3D12 190ms, Vulkan 438ms when forced with WGPU_BACKEND). Use wgpu24 for the adapters it reaches, not for speed.

Run

go run ./_example/helloworld
go run ./_example/dataset
go run ./_example/xor
go run ./_example/fizzbuzz
go run ./_example/spiral
go run ./_example/iris
go run ./_example/charrnn
go run ./_example/plasma
go run ./_example/tensor
GOEXPERIMENT=simd go run ./_example/gpt2          # downloads the GPT-2 checkpoint (~550MB) on first run
GOEXPERIMENT=simd go run ./_example/qwen -q8      # downloads Qwen2.5-0.5B-Instruct (~1GB) on first run

# The tensai command wraps the same engine as subcommands:
GOEXPERIMENT=simd go install ./cmd/tensai
tensai run -q8 "What is the capital of France?"
tensai chat -q8 -gguf model.gguf
tensai serve -q8 -addr :8080                      # OpenAI-compatible API
go run -tags wgpu ./_example/wgpu          # needs wgpu-native, see above
go run -tags wgpu ./_example/wgpu -sweep  # GPU vs CPU across sizes
go test ./...

# With the AVX2 SIMD kernel (Go 1.26+ / 1.27, amd64):
GOEXPERIMENT=simd go test ./...
GOEXPERIMENT=simd go test -bench=Dot .

The MNIST example downloads the standard IDX gzip files into _example/mnist/data when they are missing. Set MNIST_DIR to use another cache directory, and pass -model cnn for the convolutional variant (Conv2D/MaxPool2D/Dropout + AdamW); both trained variants finish by saving the model and re-scoring it after a reload. -model knn runs the no-training k-NN baseline instead — on the 5000-sample subset it scores ~91% against ~92% for the MLP and ~95% for the CNN:

go run ./_example/mnist
go run ./_example/mnist -model cnn
go run ./_example/mnist -model knn
go run ./_example/mnist -model cnn -export mnist.tflite
MNIST_DIR=/path/to/mnist go run ./_example/mnist

-export writes the trained model as a TFLite flatbuffer (the exported CNN scores identically on the LiteRT interpreter). MNIST is single-channel, so images feed the exported model unchanged; consume it from Go with go-tflite:

model := tflite.NewModelFromFile("mnist.tflite")
interpreter := tflite.NewInterpreter(model, nil)
interpreter.AllocateTensors()
copy(interpreter.GetInputTensor(0).Float32s(), image) // 28*28 floats, NHWC
interpreter.Invoke()
scores := interpreter.GetOutputTensor(0).Float32s() // 10 logits

The charrnn example trains a character-level LSTM on an embedded public-domain text, saves the parameters with SaveParamsFile, restores them into a fresh model, and generates a sample from the reloaded parameters.

The plasma example animates a demoscene-style plasma in the terminal where the plasma function is a randomly weighted network (a CPPN) evaluated for every pixel of every frame as one batch. The status line shows the per-frame network time, which makes it a live SIMD benchmark: 120x90 pixels runs at ~32 fps on the portable build and ~100 fps with GOEXPERIMENT=simd on the same machine. Try different -seed values for different effects.

Both raw IDX files and .gz variants are accepted.

SIMD Coverage

Where the AVX2 kernels apply today, and where they still could:

• Matmul (Dot/DotInto) — used by Dense, Conv2D (im2col product), knn.Classifier distances, and autograd MatMul
• ReLU / LeakyReLU forward & backward
• Sigmoid / Tanh forward & backward (vectorized polynomial exp)
• GELU forward & backward (vectorized erf)
• LayerNorm forward & backward (vector row reductions)
• Softmax / SoftmaxCrossEntropy exponentials and scaling
• Adam / AdamW parameter update
• SGD update (momentum form, same fused multiply-add loop as Adam)
• Slice add & scale primitives (bias add, Embedding gradient scatter-add)
• Transpose-free gradient matmul (DotTAInto) — Dense/Conv2D weight gradients no longer materialize input^T / im2col^T
• Remaining transposes (T/TInto, now only weight matrices and autograd) — cache-blocked 32x32 tiles
• Softmax backward row dot products (autograd) — fused AVX2 dot and Jacobian-vector accumulation
• MSE / BinaryCrossEntropy losses (BCE needs a vectorized log)
• Autograd element-wise backward passes (gradients accumulate with +=, so they need dedicated fused kernels)
• BatchNorm statistics (column-strided access needs a restructure)
• MaxPool2D window scan
• im2col / col2im gather-scatter (contiguous runs could use bulk copies)

The unchecked items are ordered roughly by expected impact; none of them show up prominently in training profiles today.

Design Notes

• All operations are batched. Inputs are MxN matrices, where M is the batch size and N is the feature dimension.
• Embedding inputs are also matrices: values must be exact integer token ids stored in Float, and the embedding vectors are flattened across the row.
• The Layer interface standardizes Forward, Backward, Params, and Grads, which keeps new layers such as convolution or dropout straightforward to add.
• Dense weights use Glorot/He-style initialization to keep early training stable.
• SoftmaxCrossEntropy subtracts the row maximum before softmax for numerical stability.

License

MIT

Author

Yasuhiro Matsumoto (a.k.a. mattn)

Documentation

Index

• func Axpy(a Float, x, y []Float)
• func Axpys(ws []Float, v, outs []Float)
• func DotInto(out, a, b *Matrix) error
• func DotTAInto(out, a, b *Matrix) error
• func DotVecs(qs, k []Float, out []Float)
• func SiluMul(gate, up []Float)
• func TInto(dst, src *Matrix) error
• type Float
  • func DotVec(a, b []Float) Float
• type Matrix
  • func Add(a, b *Matrix) (*Matrix, error)
  • func AddBias(a *Matrix, bias []Float) (*Matrix, error)
  • func Dot(a, b *Matrix) (*Matrix, error)
  • func EnsureMatrix(m *Matrix, rows, cols int) *Matrix
  • func NewMatrix(rows, cols int) *Matrix
  • func NewMatrixFromInts(rows, cols int, data []int) (*Matrix, error)
  • func NewMatrixFromSlice(rows, cols int, data []Float) (*Matrix, error)
  • func RandomMatrix(rows, cols int, rng *rand.Rand) *Matrix
  • func (m *Matrix) ArgmaxRow(r int) int
  • func (m *Matrix) At(r, c int) Float
  • func (m *Matrix) Row(r int) []Float
  • func (m *Matrix) Scale(s Float)
  • func (m *Matrix) Set(r, c int, v Float)
  • func (m *Matrix) SetRow(r int, vals []Float) error
  • func (m *Matrix) T() *Matrix
  • func (m *Matrix) Tensor() *Tensor
  • func (m *Matrix) Validate() error
• type Tensor
  • func MatMul(a, b *Tensor) (*Tensor, error)
  • func NewTensor(shape ...int) *Tensor
  • func NewTensorFromSlice(data []Float, shape ...int) (*Tensor, error)
  • func (t *Tensor) Add(o *Tensor) (*Tensor, error)
  • func (t *Tensor) At(idx ...int) Float
  • func (t *Tensor) Div(o *Tensor) (*Tensor, error)
  • func (t *Tensor) Matrix() (*Matrix, error)
  • func (t *Tensor) Mul(o *Tensor) (*Tensor, error)
  • func (t *Tensor) Reshape(shape ...int) (*Tensor, error)
  • func (t *Tensor) Scale(s Float)
  • func (t *Tensor) Set(v Float, idx ...int)
  • func (t *Tensor) Size() int
  • func (t *Tensor) Sub(o *Tensor) (*Tensor, error)
  • func (t *Tensor) Transpose(perm ...int) (*Tensor, error)
  • func (t *Tensor) Validate() error

Functions

func Axpy

func Axpy(a Float, x, y []Float)

Axpy computes y += a*x elementwise over equally long vectors — the weighted value accumulation of attention.

func Axpys

func Axpys(ws []Float, v, outs []Float)

Axpys is the grouped form of Axpy: the i-th of len(ws) rows packed contiguously in outs accumulates ws[i]*v, the shared v streamed once for up to four rows per pass — grouped-query attention's weighted value accumulation. Bit-identical to per-row Axpy.

func DotInto

func DotInto(out, a, b *Matrix) error

DotInto computes out = a * b into an existing matrix, overwriting it.

func DotTAInto

func DotTAInto(out, a, b *Matrix) error

DotTAInto computes out = a^T * b into an existing matrix, overwriting it, without materializing the transpose: a is read row by row and scattered into out with the same vector kernel Dot uses. Shapes: a is RxI, b is RxJ, out is IxJ.

func DotVecs

func DotVecs(qs, k []Float, out []Float)

DotVecs is the grouped-query form of DotVec: out[i] gets the dot of k with the i-th of len(out) query vectors packed contiguously in qs, the shared k streamed once for up to four of them per pass — the score kernel of grouped-query attention, where several query heads share one cached key row. Every result is bit-identical to the matching DotVec.

func SiluMul

func SiluMul(gate, up []Float)

SiluMul computes gate[i] = silu(gate[i]) * up[i] in place — the SwiGLU activation between a transformer block's fused gate/up projection and its down projection. The AVX2 build evaluates the sigmoid with the same polynomial exp the training kernels use, so results can differ from the portable build by a few float32 ulps.

func TInto

func TInto(dst, src *Matrix) error

TInto writes the transpose of src into dst.

Types

type Float

type Float = float32

Float is the element type of every tensor. float32 halves memory traffic versus float64 and enables the 8-lane AVX2 kernel (see dot_simd.go); its ~7 decimal digits are plenty for neural-network training.

func DotVec

func DotVec(a, b []Float) Float

DotVec returns the dot product of two equally long vectors, running on the AVX2 FMA kernel in SIMD builds — the score kernel of attention over a KV cache.

type Matrix

type Matrix struct {
	Rows int
	Cols int
	Data []Float
}

Matrix is a row-major 2D tensor of Float.

func Add

func Add(a, b *Matrix) (*Matrix, error)

Add returns a + b (element-wise). Shapes must match.

func AddBias

func AddBias(a *Matrix, bias []Float) (*Matrix, error)

AddBias adds a 1xCols bias vector to every row of a.

func Dot

func Dot(a, b *Matrix) (*Matrix, error)

Dot computes the matrix product a * b.

func EnsureMatrix

func EnsureMatrix(m *Matrix, rows, cols int) *Matrix

EnsureMatrix returns m when it already has the wanted shape, otherwise a freshly allocated matrix. The contents are unspecified; callers must overwrite (or clear) every element. Layers use this to reuse forward and backward scratch buffers between training steps.

func NewMatrix

func NewMatrix(rows, cols int) *Matrix

NewMatrix creates a matrix filled with zeros.

func NewMatrixFromInts

func NewMatrixFromInts(rows, cols int, data []int) (*Matrix, error)

NewMatrixFromInts builds a matrix from integer values, verifying each one survives the float32 conversion exactly. This is the safe way to build token-id inputs for an Embedding layer, whose ids travel as Float.

func NewMatrixFromSlice

func NewMatrixFromSlice(rows, cols int, data []Float) (*Matrix, error)

NewMatrixFromSlice creates a rows x cols matrix from row-major data.

func RandomMatrix

func RandomMatrix(rows, cols int, rng *rand.Rand) *Matrix

RandomMatrix fills a matrix with samples from a normal distribution scaled by the Glorot/Bengio gain for the given fan-in / fan-out.

func (*Matrix) ArgmaxRow

func (m *Matrix) ArgmaxRow(r int) int

ArgmaxRow returns the column index of the largest value in row r; ties go to the lowest index. Classification models emit one column per class, so this maps a row of scores to its predicted class.

func (*Matrix) At

func (m *Matrix) At(r, c int) Float

At returns the element at (r, c).

func (*Matrix) Row

func (m *Matrix) Row(r int) []Float

Row returns a copy of row r as a slice.

func (*Matrix) Scale

func (m *Matrix) Scale(s Float)

Scale multiplies every element by s, in place.

func (*Matrix) Set

func (m *Matrix) Set(r, c int, v Float)

Set sets the element at (r, c).

func (*Matrix) SetRow

func (m *Matrix) SetRow(r int, vals []Float) error

SetRow copies vals into row r.

func (*Matrix) T

func (m *Matrix) T() *Matrix

T returns the transpose of the matrix.

func (*Matrix) Tensor

func (m *Matrix) Tensor() *Tensor

Tensor returns a 2-D tensor view of the matrix sharing the same backing data.

func (*Matrix) Validate

func (m *Matrix) Validate() error

Validate returns an error if the matrix data length is inconsistent.

type Tensor

type Tensor struct {
	Shape []int
	Data  []Float
}

Tensor is an n-dimensional, contiguous, row-major array of Float — the generalization of Matrix beyond two dimensions. Element-wise arithmetic broadcasts NumPy-style: shapes are aligned at their trailing dimensions and a dimension of 1 stretches to match the other operand. MatMul multiplies stacks of matrices in one call, broadcasting the leading batch dimensions the same way.

func MatMul

func MatMul(a, b *Tensor) (*Tensor, error)

MatMul multiplies two stacks of matrices: the last two axes of each operand are the matrix dimensions and the leading axes broadcast like the element-wise ops, so a (batch..., m, k) tensor times a (batch..., k, n) tensor yields (batch..., m, n). Both operands need at least 2 axes. The per-matrix products run on the same kernel as Dot, parallelized across the batch.

func NewTensor

func NewTensor(shape ...int) *Tensor

NewTensor creates a tensor of the given shape filled with zeros.

func NewTensorFromSlice

func NewTensorFromSlice(data []Float, shape ...int) (*Tensor, error)

NewTensorFromSlice creates a tensor of the given shape from row-major data.

func (*Tensor) Add

func (t *Tensor) Add(o *Tensor) (*Tensor, error)

Add returns t + o element-wise with broadcasting.

func (*Tensor) At

func (t *Tensor) At(idx ...int) Float

At returns the element at the given multi-index.

func (*Tensor) Div

func (t *Tensor) Div(o *Tensor) (*Tensor, error)

Div returns t / o element-wise with broadcasting, with IEEE semantics for division by zero.

func (*Tensor) Matrix

func (t *Tensor) Matrix() (*Matrix, error)

Matrix returns a matrix view of a 2-D tensor sharing the same backing data.

func (*Tensor) Mul

func (t *Tensor) Mul(o *Tensor) (*Tensor, error)

Mul returns t * o element-wise with broadcasting.

func (*Tensor) Reshape

func (t *Tensor) Reshape(shape ...int) (*Tensor, error)

Reshape returns a tensor with a new shape sharing the same backing data. One dimension may be -1 and is inferred from the element count.

func (*Tensor) Scale

func (t *Tensor) Scale(s Float)

Scale multiplies every element by s, in place.

func (*Tensor) Set

func (t *Tensor) Set(v Float, idx ...int)

Set sets the element at the given multi-index.

func (*Tensor) Size

func (t *Tensor) Size() int

Size returns the total number of elements.

func (*Tensor) Sub

func (t *Tensor) Sub(o *Tensor) (*Tensor, error)

Sub returns t - o element-wise with broadcasting.

func (*Tensor) Transpose

func (t *Tensor) Transpose(perm ...int) (*Tensor, error)

Transpose returns a copy of the tensor with its axes permuted; perm must list every axis exactly once. With no arguments it swaps the last two axes — the matrix transpose of every matrix in the stack — matching Matrix.T for 2-D tensors.

func (*Tensor) Validate

func (t *Tensor) Validate() error

Validate returns an error if the tensor shape or data length is inconsistent.