tensai

package module
v0.0.2 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Aug 27, 2026 License: MIT Imports: 10 Imported by: 0

README

tensai

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).

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 OpenGPU() 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 - QuantizeMatrix / QuantizeMatrix4 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 - RNNCell, LSTMCell, and single-head 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)
tensor.go           Matrix and vector operations
ndtensor.go         N-d Tensor: broadcasting element-wise ops, batched MatMul
wgpu.go             WebGPU MatMul backend via purego + wgpu-native (build tag wgpu)
dot_simd.go         AVX2 matmul kernel (GOEXPERIMENT=simd, amd64)
dot_generic.go      Portable matmul kernel (all other builds)
kernels.go          Scalar bodies of the element-wise kernels
mathvec_simd.go     AVX2 element-wise kernels incl. vectorized exp
mathvec_generic.go  Portable element-wise kernel dispatchers
layer.go            Layer interface plus Dense and activations
conv.go             Conv2D and MaxPool2D layers
batchnorm.go        BatchNorm layer
dropout.go          Dropout layer
loss.go             Loss functions (MSE, SoftmaxCrossEntropy, BCE)
optimizer.go        Optimizers (SGD, Adam, AdamW)
model.go            Sequential model and training loop
autograd.go         Reverse-mode automatic differentiation (Node graph)
rnn.go              RNNCell / LSTMCell / SelfAttention on the autograd engine
serialize.go        Model parameter save/load (JSON)
tensai_test.go      Unit tests plus XOR convergence test
features_test.go    Gradient checks and tests for the newer layers
_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
model := tensai.NewSequential()
model.Add(tensai.NewDense(8))
model.Add(&tensai.Tanh{})
model.Add(tensai.NewDense(1))
model.Add(&tensai.Sigmoid{})

model.Compile(2, tensai.MeanSquaredError{}, tensai.NewAdam(0.05))
model.Fit(inputs, targets, 5000)

pred, _ := model.Predict(inputs)
Classification: softmax + cross-entropy
model := tensai.NewSequential()
model.Add(tensai.NewDense(8))
model.Add(&tensai.ReLU{})
model.Add(tensai.NewDense(2)) // output width = number of classes

model.Compile(2, tensai.SoftmaxCrossEntropy{}, tensai.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, _ := tensai.NewDataset(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 := model.FitStep(in, tgt)
		return err
	})
}
Convolution, regularization, and saving
model := tensai.NewSequential()
model.Add(tensai.NewConv2D(28, 28, 1, 8, 3, 1, 1)) // inH, inW, inC, outC, kernel, stride, pad
model.Add(&tensai.ReLU{})
model.Add(tensai.NewMaxPool2D(28, 28, 8, 2))
model.Add(tensai.NewDense(64))
model.Add(tensai.NewBatchNorm())
model.Add(tensai.NewLeakyReLU(0.01))
model.Add(tensai.NewDropout(0.3))
model.Add(tensai.NewDense(10))

model.Compile(28*28, tensai.SoftmaxCrossEntropy{}, tensai.NewAdamW(0.001, 0.01))
model.Fit(inputs, targets, 10)

model.SaveFile("model.json")
// Later: build + Compile the same architecture, then
model.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 := tensai.Param(tensai.RandomMatrix(2, 8, rng))
b1 := tensai.Param(tensai.NewMatrix(1, 8))
w2 := tensai.Param(tensai.RandomMatrix(8, 1, rng))
trainer := tensai.NewTrainer(tensai.NewAdam(0.05), w1, b1, w2)

for step := 0; step < 2000; step++ {
	loss := tensai.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 tensai.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

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

cell := tensai.NewLSTMCell(inSize, hidden, rng)
wOut := tensai.Param(tensai.RandomMatrix(hidden, numClasses, rng))
bOut := tensai.Param(tensai.NewMatrix(1, numClasses))
trainer := tensai.NewTrainer(tensai.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(tensai.Input(x), h, c)
	}
	logits := h.MatMul(wOut).AddRow(bOut)
	trainer.Step(logits.SoftmaxCELoss(labels))
}

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 tensai.Attention(q, k, v) form is also exposed.

Autograd parameters are saved and restored positionally with tensai.SaveParamsFile("cell.json", cell.Params()...) / tensai.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:

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

On machines with both an integrated and a discrete GPU, pass a preference: tensai.OpenGPU(tensai.GPULowPower) steers to the iGPU, tensai.GPUHighPerformance 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, _ := gpu.Upload(w)              // weight uploaded once
defer gw.Free()                     // GPU memory is not garbage collected
gx, _ := gpu.Upload(x)
h, _ := gx.MatMul(gw)               // chain freely; nothing touches the host
out, _ := h.MatMul(gw2)
result, _ := out.Download()         // one readback at the end

gpu.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 GPUQMatrix.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 GPUTensor.MatMul keeps inputs and intermediates resident, so only the final result needs to cross the bus. Without the build tag OpenGPU 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 OpenGPU 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 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
  • 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)
  • SGD update

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

Constants

This section is empty.

Variables

This section is empty.

Functions

func Axpy added in v0.0.2 

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 added in v0.0.2 

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 added in v0.0.2 

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 LoadParams 

func LoadParams(r io.Reader, params ...*Node) error

LoadParams restores parameter values saved by SaveParams. The parameters must be passed in the same order and have the same shapes as when saved.

func LoadParamsFile 

func LoadParamsFile(path string, params ...*Node) error

LoadParamsFile restores autograd parameters from a file written by SaveParamsFile.

func MXFP4Value added in v0.0.2 

func MXFP4Value(code uint8) int8

MXFP4Value returns an FP4 code's expanded integer value (twice the FP4 value); callers building ColSum64 sum these.

func PackScaleMin added in v0.0.2 

func PackScaleMin(scale, min Float) uint32

PackScaleMin rounds a min-form group's scale and min to bfloat16 and packs them into one ScaleMin entry (scale low, min high).

func SaveParams 

func SaveParams(w io.Writer, params ...*Node) error

SaveParams writes the values of autograd parameters as JSON, in the given order. Pass the same parameter list a Trainer uses, e.g. SaveParams(w, cell.Params()...).

func SaveParamsFile 

func SaveParamsFile(path string, params ...*Node) error

SaveParamsFile writes autograd parameters to a JSON file.

func SiluMul added in v0.0.2 

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.

func ZeroGrads 

func ZeroGrads(nodes ...*Node)

ZeroGrads clears the gradients of the given nodes. Call it between training steps: Backward accumulates.

Types

type Adam 

type Adam struct {
	LR          Float
	Beta1       Float
	Beta2       Float
	Eps         Float
	WeightDecay Float
	// contains filtered or unexported fields
}

Adam optimizer. With WeightDecay > 0 it becomes AdamW: decay is decoupled from the gradient update and applied to weights only (never to biases).

func NewAdam 

func NewAdam(lr Float) *Adam

NewAdam returns an Adam optimizer with standard defaults.

func NewAdamW 

func NewAdamW(lr, weightDecay Float) *Adam

NewAdamW returns an Adam optimizer with decoupled weight decay (AdamW).

func (*Adam) Name 

func (a *Adam) Name() string

Name returns "adam".

func (*Adam) NewLayer 

func (a *Adam) NewLayer() int

NewLayer registers a layer and returns its index.

func (*Adam) Step 

func (a *Adam) Step(idx int, weights, gradW *Matrix, bias, gradB []Float)

Step updates parameters with the Adam rule.

func (*Adam) String 

func (a *Adam) String() string

String returns a human-readable description of an optimizer's config.

type BatchNorm 

type BatchNorm struct {
	Momentum Float // running-stats decay, default 0.9
	Eps      Float // numerical stability, default 1e-5
	// contains filtered or unexported fields
}

BatchNorm normalizes each feature column over the batch, then applies a learned scale (gamma) and shift (beta). During training it normalizes with batch statistics and maintains running estimates; during inference it uses the running estimates.

gamma is exposed as the layer's weights (a 1xC matrix) and beta as its bias, so optimizers update them like any other parameters.

func NewBatchNorm 

func NewBatchNorm() *BatchNorm

NewBatchNorm returns a BatchNorm layer with standard defaults.

func (*BatchNorm) Backward 

func (b *BatchNorm) Backward(gradOutput *Matrix) (*Matrix, error)

func (*BatchNorm) Forward 

func (b *BatchNorm) Forward(input *Matrix) (*Matrix, error)

func (*BatchNorm) Grads 

func (b *BatchNorm) Grads() (*Matrix, []Float)

func (*BatchNorm) Init 

func (b *BatchNorm) Init(inputCols int, _ *rand.Rand) (int, error)

func (*BatchNorm) Params 

func (b *BatchNorm) Params() (*Matrix, []Float)

func (*BatchNorm) RunningStats 

func (b *BatchNorm) RunningStats() (mean, variance []Float)

RunningStats returns the running mean and variance estimates used at inference time, for exporters.

func (*BatchNorm) SetParams 

func (b *BatchNorm) SetParams(weights *Matrix, bias []Float) error

type BinaryCrossEntropy 

type BinaryCrossEntropy struct{}

BinaryCrossEntropy computes the average binary cross-entropy between predicted probabilities in (0,1) and 0/1 targets of the same shape. Pair it with a Sigmoid output layer.

func (BinaryCrossEntropy) Loss 

func (BinaryCrossEntropy) Loss(pred, target *Matrix) (Float, *Matrix, error)

Loss returns the average binary cross-entropy and its gradient.

func (BinaryCrossEntropy) LossInto 

func (BinaryCrossEntropy) LossInto(pred, target, grad *Matrix) (Float, error)

LossInto writes the BCE gradient into grad and returns the average loss.

func (BinaryCrossEntropy) Name 

func (BinaryCrossEntropy) Name() string

Name returns "bce".

type Conv2D 

type Conv2D struct {
	// contains filtered or unexported fields
}

Conv2D is a 2D convolution layer. Because the framework moves data as flat MxN matrices, each sample row must be laid out channel-major: index = (channel*height + y)*width + x. The output uses the same layout.

The convolution is computed as a matrix product over an im2col expansion, so it reuses the tuned Dot kernel. Weights are stored as an (inC*kernel*kernel) x outC matrix.

func NewConv2D 

func NewConv2D(inH, inW, inC, outC, kernel, stride, pad int) *Conv2D

NewConv2D returns a convolution layer for inH x inW inputs with inC channels, producing outC channels with a square kernel.

func (*Conv2D) Backward 

func (c *Conv2D) Backward(gradOutput *Matrix) (*Matrix, error)

func (*Conv2D) Forward 

func (c *Conv2D) Forward(input *Matrix) (*Matrix, error)

func (*Conv2D) Grads 

func (c *Conv2D) Grads() (*Matrix, []Float)

func (*Conv2D) Init 

func (c *Conv2D) Init(inputCols int, rng *rand.Rand) (int, error)

func (*Conv2D) Params 

func (c *Conv2D) Params() (*Matrix, []Float)

func (*Conv2D) SetParams 

func (c *Conv2D) SetParams(weights *Matrix, bias []Float) error

func (*Conv2D) Shape 

func (c *Conv2D) Shape() (inH, inW, inC, outC, kernel, stride, pad int)

Shape reports the layer's spatial configuration, for exporters.

type Dataset added in v0.0.2 

type Dataset struct {
	Inputs  *Matrix
	Targets *Matrix
}

Dataset pairs an input matrix with its target matrix, row-aligned, and provides the usual training-data plumbing: shuffling, train/test splitting, mini-batch iteration, and standardization.

func NewDataset added in v0.0.2 

func NewDataset(inputs, targets *Matrix) (*Dataset, error)

NewDataset wraps inputs and targets after checking row alignment.

func (*Dataset) Batches added in v0.0.2 

func (d *Dataset) Batches(size int, rng *rand.Rand, fn func(inputs, targets *Matrix) error) error

Batches invokes fn once per mini-batch of exactly size samples, copying rows into buffers that are reused between calls (do not retain them). With a non-nil rng the visit order is reshuffled; trailing samples that do not fill a batch are skipped, matching common epoch loops.

func (*Dataset) Len added in v0.0.2 

func (d *Dataset) Len() int

Len returns the number of samples.

func (*Dataset) Shuffle added in v0.0.2 

func (d *Dataset) Shuffle(rng *rand.Rand)

Shuffle permutes the samples in place, keeping input and target rows paired.

func (*Dataset) Split added in v0.0.2 

func (d *Dataset) Split(testFraction float64) (train, test *Dataset, err error)

Split divides the dataset into a training and a test set, with testFraction (0 < f < 1) of the samples going to the test set. The two halves are views sharing the underlying data — no rows are copied. Shuffle first when the data is ordered.

func (*Dataset) Standardize added in v0.0.2 

func (d *Dataset) Standardize() (mean, std []Float)

Standardize scales every input column to zero mean and unit variance in place and returns the per-column statistics, for applying the same transform to other data with StandardizeWith. Constant columns keep a standard deviation of 1 so they pass through unchanged.

func (*Dataset) StandardizeWith added in v0.0.2 

func (d *Dataset) StandardizeWith(mean, std []Float)

StandardizeWith applies previously computed statistics to the inputs in place (e.g. training-set statistics to a test set).

type Dense 

type Dense struct {
	// contains filtered or unexported fields
}

Dense is a fully-connected layer: y = x*W + b.

func NewDense 

func NewDense(outCols int) *Dense

NewDense returns a Dense layer with the given output size.

func (*Dense) Backward 

func (d *Dense) Backward(gradOutput *Matrix) (*Matrix, error)

func (*Dense) Forward 

func (d *Dense) Forward(input *Matrix) (*Matrix, error)

func (*Dense) Grads 

func (d *Dense) Grads() (*Matrix, []Float)

func (*Dense) Init 

func (d *Dense) Init(inputCols int, rng *rand.Rand) (int, error)

func (*Dense) Params 

func (d *Dense) Params() (*Matrix, []Float)

func (*Dense) SetParams 

func (d *Dense) SetParams(weights *Matrix, bias []Float) error

type Dropout 

type Dropout struct {
	Rate Float
	// contains filtered or unexported fields
}

Dropout randomly zeroes elements during training with probability Rate and scales the survivors by 1/(1-Rate) ("inverted dropout"), so inference is a plain pass-through with no rescaling.

func NewDropout 

func NewDropout(rate Float) *Dropout

NewDropout returns a Dropout layer that drops the given fraction of activations during training. Rate must be in [0, 1).

func (*Dropout) Backward 

func (d *Dropout) Backward(gradOutput *Matrix) (*Matrix, error)

func (*Dropout) Forward 

func (d *Dropout) Forward(input *Matrix) (*Matrix, error)

func (*Dropout) Grads 

func (d *Dropout) Grads() (*Matrix, []Float)

func (*Dropout) Init 

func (d *Dropout) Init(inputCols int, rng *rand.Rand) (int, error)

func (*Dropout) Params 

func (d *Dropout) Params() (*Matrix, []Float)

func (*Dropout) SetParams 

func (d *Dropout) SetParams(*Matrix, []Float) error

type Embedding 

type Embedding struct {
	// contains filtered or unexported fields
}

Embedding looks up a learned vector for each token id in the input row and concatenates the vectors across columns.

func NewEmbedding 

func NewEmbedding(vocabSize, dim int) *Embedding

NewEmbedding returns a trainable embedding table of shape vocabSize x dim.

func (*Embedding) Backward 

func (e *Embedding) Backward(gradOutput *Matrix) (*Matrix, error)

func (*Embedding) Forward 

func (e *Embedding) Forward(input *Matrix) (*Matrix, error)

func (*Embedding) Grads 

func (e *Embedding) Grads() (*Matrix, []Float)

func (*Embedding) Init 

func (e *Embedding) Init(inputCols int, rng *rand.Rand) (int, error)

func (*Embedding) Params 

func (e *Embedding) Params() (*Matrix, []Float)

func (*Embedding) SetParams 

func (e *Embedding) SetParams(weights *Matrix, bias []Float) error

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 added in v0.0.2 

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.

func MXFP4Scale added in v0.0.2 

func MXFP4Scale(e uint8) Float

MXFP4Scale converts an E8M0 exponent byte to the matrix's Scale entry: half of 2^(e-127), matching the doubled integer grid.

func UnpackScaleMin added in v0.0.2 

func UnpackScaleMin(u uint32) (scale, min Float)

UnpackScaleMin is the inverse of PackScaleMin.

type GELU 

type GELU struct {
	// contains filtered or unexported fields
}

GELU activation: f(x) = 0.5*x*(1+erf(x/sqrt(2))).

func (*GELU) Backward 

func (g *GELU) Backward(gradOutput *Matrix) (*Matrix, error)

func (*GELU) Forward 

func (g *GELU) Forward(input *Matrix) (*Matrix, error)

func (*GELU) Grads 

func (a *GELU) Grads() (*Matrix, []Float)

func (*GELU) Init 

func (g *GELU) Init(inputCols int, _ *rand.Rand) (int, error)

func (*GELU) Params 

func (a *GELU) Params() (*Matrix, []Float)

func (*GELU) SetParams 

func (a *GELU) SetParams(*Matrix, []Float) error

type GPU added in v0.0.2 

type GPU struct{}

GPU is the WebGPU compute backend. This build has it disabled; build with -tags wgpu (linux, darwin, or windows) and see wgpu.go for the runtime requirements.

func OpenGPU added in v0.0.2 

func OpenGPU(power ...GPUPower) (*GPU, error)

OpenGPU always fails in builds without the wgpu tag.

func (*GPU) BeginBatch added in v0.0.2 

func (g *GPU) BeginBatch() error

BeginBatch always fails in builds without the wgpu tag.

func (*GPU) Close added in v0.0.2 

func (g *GPU) Close()

Close is a no-op in builds without the wgpu tag.

func (*GPU) Flush added in v0.0.2 

func (g *GPU) Flush() error

Flush is a no-op in builds without the wgpu tag.

func (*GPU) HasF16 added in v0.0.2 

func (g *GPU) HasF16() bool

HasF16 reports false in builds without the wgpu tag.

func (*GPU) MatMul added in v0.0.2 

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

MatMul always fails in builds without the wgpu tag.

func (*GPU) Name added in v0.0.2 

func (g *GPU) Name() string

Name always returns "" in builds without the wgpu tag.

func (*GPU) NewF16Tensor added in v0.0.2 

func (g *GPU) NewF16Tensor(shape ...int) (*GPUTensor, error)

NewF16Tensor always fails in builds without the wgpu tag.

func (*GPU) StorageLimit added in v0.0.2 

func (g *GPU) StorageLimit() uint64

StorageLimit returns 0 in builds without the wgpu tag.

func (*GPU) Upload added in v0.0.2 

func (g *GPU) Upload(t *Tensor) (*GPUTensor, error)

Upload always fails in builds without the wgpu tag.

func (*GPU) UploadQ4 added in v0.0.2 

func (g *GPU) UploadQ4(q *Q4Matrix) (*GPUQ4Matrix, error)

UploadQ4 always fails in builds without the wgpu tag.

func (*GPU) UploadQ8 added in v0.0.2 

func (g *GPU) UploadQ8(q *QMatrix) (*GPUQMatrix, error)

UploadQ8 always fails in builds without the wgpu tag.

type GPUPower added in v0.0.2 

type GPUPower uint32

GPUPower tells OpenGPU which adapter to prefer; unused in builds without the wgpu tag. 

const (
	GPUDefault         GPUPower = 0
	GPULowPower        GPUPower = 1
	GPUHighPerformance GPUPower = 2
)

type GPUQ4Matrix added in v0.0.2 

type GPUQ4Matrix struct{}

GPUQ4Matrix is a GPU-resident int4 weight matrix. This build has it disabled; build with -tags wgpu or -tags wgpu24 to enable it.

func (*GPUQ4Matrix) Free added in v0.0.2 

func (q *GPUQ4Matrix) Free()

Free is a no-op in builds without the wgpu tag.

func (*GPUQ4Matrix) MatMul added in v0.0.2 

func (q *GPUQ4Matrix) MatMul(x *GPUTensor) (*GPUTensor, error)

MatMul always fails in builds without the wgpu tag.

func (*GPUQ4Matrix) MatMulOpts added in v0.0.2 

func (q *GPUQ4Matrix) MatMulOpts(x, bias, dst *GPUTensor) (*GPUTensor, error)

MatMulOpts always fails in builds without the wgpu tag.

func (*GPUQ4Matrix) Shape added in v0.0.2 

func (q *GPUQ4Matrix) Shape() (int, int)

Shape returns zeros in builds without the wgpu tag.

type GPUQMatrix added in v0.0.2 

type GPUQMatrix struct{}

GPUQMatrix is a GPU-resident int8 weight matrix. This build has it disabled; build with -tags wgpu or -tags wgpu24 to enable it.

func (*GPUQMatrix) Free added in v0.0.2 

func (q *GPUQMatrix) Free()

Free is a no-op in builds without the wgpu tag.

func (*GPUQMatrix) MatMul added in v0.0.2 

func (q *GPUQMatrix) MatMul(x *GPUTensor) (*GPUTensor, error)

MatMul always fails in builds without the wgpu tag.

func (*GPUQMatrix) MatMulOpts added in v0.0.2 

func (q *GPUQMatrix) MatMulOpts(x, bias, dst *GPUTensor) (*GPUTensor, error)

MatMulOpts always fails in builds without the wgpu tag.

func (*GPUQMatrix) Shape added in v0.0.2 

func (q *GPUQMatrix) Shape() (int, int)

Shape returns zeros in builds without the wgpu tag.

type GPUTensor added in v0.0.2 

type GPUTensor struct{}

GPUTensor is a GPU-resident tensor. This build has it disabled; build with -tags wgpu or -tags wgpu24 to enable it.

func (*GPUTensor) Add added in v0.0.2 

func (t *GPUTensor) Add(o *GPUTensor) error

Add always fails in builds without the wgpu tag.

func (*GPUTensor) Attention added in v0.0.2 

func (q *GPUTensor) Attention(k, v *GPUTensor) (*GPUTensor, error)

Attention always fails in builds without the wgpu tag.

func (*GPUTensor) CausalAttention added in v0.0.2 

func (q *GPUTensor) CausalAttention(k, v *GPUTensor) (*GPUTensor, error)

CausalAttention always fails in builds without the wgpu tag.

func (*GPUTensor) CausalMultiHeadAttention added in v0.0.2 

func (q *GPUTensor) CausalMultiHeadAttention(k, v *GPUTensor, heads int) (*GPUTensor, error)

CausalMultiHeadAttention always fails in builds without the wgpu tag.

func (*GPUTensor) CopyRowsInto added in v0.0.2 

func (t *GPUTensor) CopyRowsInto(dst *GPUTensor, off int) error

CopyRowsInto always fails in builds without the wgpu tag.

func (*GPUTensor) Download added in v0.0.2 

func (t *GPUTensor) Download() (*Tensor, error)

Download always fails in builds without the wgpu tag.

func (*GPUTensor) DownloadRange added in v0.0.2 

func (t *GPUTensor) DownloadRange(off, n int) (*Tensor, error)

DownloadRange always fails in builds without the wgpu tag.

func (*GPUTensor) Free added in v0.0.2 

func (t *GPUTensor) Free()

Free is a no-op in builds without the wgpu tag.

func (*GPUTensor) GeluMul added in v0.0.2 

func (t *GPUTensor) GeluMul(o *GPUTensor) error

GeluMul always fails in builds without the wgpu tag.

func (*GPUTensor) GroupedCausalAttention added in v0.0.2 

func (q *GPUTensor) GroupedCausalAttention(k, v *GPUTensor, heads, kvHeads, seqKV, window int) (*GPUTensor, error)

GroupedCausalAttention always fails in builds without the wgpu tag.

func (*GPUTensor) MatMul added in v0.0.2 

func (t *GPUTensor) MatMul(o *GPUTensor) (*GPUTensor, error)

MatMul always fails in builds without the wgpu tag.

func (*GPUTensor) MatMulT added in v0.0.2 

func (t *GPUTensor) MatMulT(o *GPUTensor) (*GPUTensor, error)

MatMulT always fails in builds without the wgpu tag.

func (*GPUTensor) MultiHeadAttention added in v0.0.2 

func (q *GPUTensor) MultiHeadAttention(k, v *GPUTensor, heads int) (*GPUTensor, error)

MultiHeadAttention always fails in builds without the wgpu tag.

func (*GPUTensor) RMSNorm added in v0.0.2 

func (t *GPUTensor) RMSNorm(w *GPUTensor, eps float64) (*GPUTensor, error)

RMSNorm always fails in builds without the wgpu tag.

func (*GPUTensor) RMSNormEach added in v0.0.2 

func (t *GPUTensor) RMSNormEach(w *GPUTensor, eps float64) (*GPUTensor, error)

RMSNormEach always fails in builds without the wgpu tag.

func (*GPUTensor) RoPE added in v0.0.2 

func (t *GPUTensor) RoPE(headSz, pos0 int, theta float64) error

RoPE always fails in builds without the wgpu tag.

func (*GPUTensor) Scale added in v0.0.2 

func (t *GPUTensor) Scale(s Float) error

Scale always fails in builds without the wgpu tag.

func (*GPUTensor) Shape added in v0.0.2 

func (t *GPUTensor) Shape() []int

Shape returns nil in builds without the wgpu tag.

func (*GPUTensor) SiluMul added in v0.0.2 

func (t *GPUTensor) SiluMul(o *GPUTensor) error

SiluMul always fails in builds without the wgpu tag.

func (*GPUTensor) Size added in v0.0.2 

func (t *GPUTensor) Size() int

Size returns 0 in builds without the wgpu tag.

func (*GPUTensor) Softmax added in v0.0.2 

func (t *GPUTensor) Softmax() (*GPUTensor, error)

Softmax always fails in builds without the wgpu tag.

type KNN 

type KNN struct {
	K int
	// contains filtered or unexported fields
}

KNN is a k-nearest-neighbors classifier. It is a lazy learner: Fit just stores the training data, and Predict ranks neighbors by squared Euclidean distance. The distance computation is reduced to one matrix product per chunk (||a-b||^2 = ||a||^2 + ||b||^2 - 2*a.b), so it runs on the same tuned Dot kernel as the neural networks.

func NewKNN 

func NewKNN(k int) *KNN

NewKNN returns a classifier that votes among the k nearest neighbors.

func (*KNN) Fit 

func (k *KNN) Fit(inputs, targets *Matrix) error

Fit stores the training set. Targets must be an Mx1 matrix of class indices, as with SoftmaxCrossEntropy.

func (*KNN) Predict 

func (k *KNN) Predict(inputs *Matrix) (*Matrix, error)

Predict returns an (inputs.Rows x classes) matrix of neighbor vote counts; take the argmax per row for the predicted class. Row sums equal K.

type LSTMCell 

type LSTMCell struct {
	// One (Wx, Wh, B) triple per gate: forget, input, output, candidate.
	Wxf, Whf, Bf *Node
	Wxi, Whi, Bi *Node
	Wxo, Who, Bo *Node
	Wxg, Whg, Bg *Node
}

LSTMCell is a long short-term memory cell with forget/input/output gates.

func NewLSTMCell 

func NewLSTMCell(inSize, hidden int, rng *rand.Rand) *LSTMCell

NewLSTMCell returns a randomly initialized LSTM cell. Forget-gate biases start at 1 so early training defaults to remembering.

func (*LSTMCell) InitState 

func (c *LSTMCell) InitState(batch int) (*Node, *Node)

InitState returns zero (hidden, cell) states for the given batch size.

func (*LSTMCell) Params 

func (c *LSTMCell) Params() []*Node

Params returns the cell's trainable parameters, for NewTrainer.

func (*LSTMCell) Step 

func (c *LSTMCell) Step(x, h, cell *Node) (*Node, *Node)

Step consumes one time step with the previous (hidden, cell) state and returns the next (hidden, cell) state.

type Layer 

type Layer interface {
	// Init configures parameters using the given RNG and input width.
	Init(inputCols int, rng *rand.Rand) (outputCols int, err error)

	// Forward computes activations given the input batch.
	Forward(input *Matrix) (*Matrix, error)

	// Backward computes the gradient with respect to the layer input,
	// given the gradient with respect to the layer output.
	Backward(gradOutput *Matrix) (*Matrix, error)

	// Grads returns the parameter gradients accumulated during the last
	// backward pass, in the order [weights, bias]. Layers without
	// parameters return nil.
	Grads() (*Matrix, []Float)

	// Params returns the current parameters [weights, bias].
	Params() (*Matrix, []Float)

	// SetParams replaces the parameters [weights, bias].
	SetParams(weights *Matrix, bias []Float) error
}

Layer is a single differentiable stage of a Sequential model. Forward and Backward are batched: inputs/outputs are MxN matrices where M is the batch size and N is the feature dimension.

type LayerNorm 

type LayerNorm struct {
	// contains filtered or unexported fields
}

LayerNorm normalizes each row over its feature dimension and applies a learnable affine transform.

func NewLayerNorm 

func NewLayerNorm() *LayerNorm

NewLayerNorm returns a LayerNorm with the default epsilon.

func (*LayerNorm) Backward 

func (l *LayerNorm) Backward(gradOutput *Matrix) (*Matrix, error)

func (*LayerNorm) Forward 

func (l *LayerNorm) Forward(input *Matrix) (*Matrix, error)

func (*LayerNorm) Grads 

func (l *LayerNorm) Grads() (*Matrix, []Float)

func (*LayerNorm) Init 

func (l *LayerNorm) Init(inputCols int, _ *rand.Rand) (int, error)

func (*LayerNorm) Params 

func (l *LayerNorm) Params() (*Matrix, []Float)

func (*LayerNorm) SetParams 

func (l *LayerNorm) SetParams(weights *Matrix, bias []Float) error

type LeakyReLU 

type LeakyReLU struct {
	Alpha Float
	// contains filtered or unexported fields
}

LeakyReLU activation: f(x) = x for x > 0, alpha*x otherwise.

func NewLeakyReLU 

func NewLeakyReLU(alpha Float) *LeakyReLU

NewLeakyReLU returns a LeakyReLU with the given negative-side slope.

func (*LeakyReLU) Backward 

func (l *LeakyReLU) Backward(gradOutput *Matrix) (*Matrix, error)

func (*LeakyReLU) Forward 

func (l *LeakyReLU) Forward(input *Matrix) (*Matrix, error)

func (*LeakyReLU) Grads 

func (a *LeakyReLU) Grads() (*Matrix, []Float)

func (*LeakyReLU) Init 

func (l *LeakyReLU) Init(inputCols int, _ *rand.Rand) (int, error)

func (*LeakyReLU) Params 

func (a *LeakyReLU) Params() (*Matrix, []Float)

func (*LeakyReLU) SetParams 

func (a *LeakyReLU) SetParams(*Matrix, []Float) error

type Loss 

type Loss interface {
	// Loss returns the average loss and the per-element gradient dL/dpred.
	Loss(pred, target *Matrix) (Float, *Matrix, error)
	// Name is a short identifier for logging.
	Name() string
}

Loss is a differentiable loss function operating on a prediction batch and a target batch of the same shape. It returns the scalar loss and the gradient of the loss with respect to the predictions.

type MXFP4Matrix added in v0.0.2 

type MXFP4Matrix struct {
	Rows, Cols int
	Q          []uint8 // FP4 codes, tiled quad nibbles (Index)
	Scale      []Float // per (32-row group, column), tile-major (TableIndex)
	ColSum64   []int32 // 64 * sum of a group's expanded codes, tile-major
}

MXFP4Matrix is a weight matrix in microscaling FP4: 32-row groups per column share one power-of-two E8M0 factor, and each weight is a 4-bit FP4 code (E2M1: 0, ±0.5, ±1, ±1.5, ±2, ±3, ±4, ±6 before scaling) — the format gpt-oss ships its expert weights in. Codes store in Q4Matrix's tiled quad-nibble layout; the kernels expand each code to twice its FP4 value on the integer grid {0, ±1..±4, ±6, ±8, ±12} with a 16-entry table lookup and run the grouped-int8 multiply-add chain, Scale carrying half the E8M0 factor so the products stay exact.

func NewMXFP4Matrix added in v0.0.2 

func NewMXFP4Matrix(rows, cols int) *MXFP4Matrix

NewMXFP4Matrix allocates the layout for rows x cols; the caller fills Q via Index and the tile-major tables via TableIndex.

func (*MXFP4Matrix) Index added in v0.0.2 

func (q *MXFP4Matrix) Index(i, j int) int

Index returns the position in Q of the byte carrying column j of rows i and i+1 (low and high nibble; shift by 4*(i%2)).

func (*MXFP4Matrix) MatMul added in v0.0.2 

func (q *MXFP4Matrix) MatMul(x, out *Matrix) error

MatMul computes out = x @ Q for a batch of activation rows.

func (*MXFP4Matrix) MatVec added in v0.0.2 

func (q *MXFP4Matrix) MatVec(x, out []Float) error

MatVec computes out = x @ Q for a single activation row: len(x) must be Rows and len(out) Cols.

func (*MXFP4Matrix) TableIndex added in v0.0.2 

func (q *MXFP4Matrix) TableIndex(g, j int) int

TableIndex returns the position in Scale and ColSum64 of group g, column j.

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 NewMatrix 

func NewMatrix(rows, cols int) *Matrix

NewMatrix creates a matrix filled with zeros.

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) 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 added in v0.0.2 

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 MaxPool2D 

type MaxPool2D struct {
	// contains filtered or unexported fields
}

MaxPool2D downsamples each channel by taking the maximum over non-overlapping size x size windows. It expects the same channel-major layout as Conv2D.

func NewMaxPool2D 

func NewMaxPool2D(inH, inW, channels, size int) *MaxPool2D

NewMaxPool2D returns a max-pooling layer with stride equal to size.

func (*MaxPool2D) Backward 

func (p *MaxPool2D) Backward(gradOutput *Matrix) (*Matrix, error)

func (*MaxPool2D) Forward 

func (p *MaxPool2D) Forward(input *Matrix) (*Matrix, error)

func (*MaxPool2D) Grads 

func (p *MaxPool2D) Grads() (*Matrix, []Float)

func (*MaxPool2D) Init 

func (p *MaxPool2D) Init(inputCols int, _ *rand.Rand) (int, error)

func (*MaxPool2D) Params 

func (p *MaxPool2D) Params() (*Matrix, []Float)

func (*MaxPool2D) SetParams 

func (p *MaxPool2D) SetParams(*Matrix, []Float) error

func (*MaxPool2D) Shape 

func (p *MaxPool2D) Shape() (inH, inW, channels, size int)

Shape reports the layer's spatial configuration, for exporters.

type MeanSquaredError 

type MeanSquaredError struct{}

MeanSquaredError computes the average of squared differences.

func (MeanSquaredError) Loss 

func (MeanSquaredError) Loss(pred, target *Matrix) (Float, *Matrix, error)

Loss returns the mean squared error and its gradient.

func (MeanSquaredError) LossInto 

func (MeanSquaredError) LossInto(pred, target, grad *Matrix) (Float, error)

LossInto writes the MSE gradient into grad and returns the average loss.

func (MeanSquaredError) Name 

func (MeanSquaredError) Name() string

Name returns "mse".

type Model 

type Model interface {
	Predict(input *Matrix) (*Matrix, error)
}

Model is a trained, ready-to-predict network.

type Node 

type Node struct {
	Value *Matrix
	Grad  *Matrix
	// contains filtered or unexported fields
}

Node is a matrix-valued node in a dynamically built computation graph for reverse-mode automatic differentiation. Build the forward computation by chaining operations, then call Backward on the (scalar) result to fill Grad on every Param node that contributed to it.

Unlike the Layer API, shape mismatches panic: graph construction errors are programming errors, and error returns would make chaining unusable. 

w := tensai.Param(tensai.RandomMatrix(2, 8, rng))
b := tensai.Param(tensai.NewMatrix(1, 8))
loss := tensai.Input(x).MatMul(w).AddRow(b).ReLU().MSELoss(y)
loss.Backward()
// w.Grad and b.Grad now hold dLoss/dw and dLoss/db.

func Attention 

func Attention(q, k, v *Node) *Node

Attention computes scaled dot-product attention softmax(q*k^T/sqrt(d))*v for a single sequence, where q, k, v are (seqLen x d) nodes.

func Input 

func Input(m *Matrix) *Node

Input wraps a matrix as a constant graph leaf. No gradient is computed for it.

func Param 

func Param(m *Matrix) *Node

Param wraps a matrix as a trainable graph leaf. Backward accumulates its gradient into Grad.

func (*Node) Add 

func (n *Node) Add(o *Node) *Node

Add returns the element-wise sum n + o. Shapes must match.

func (*Node) AddRow 

func (n *Node) AddRow(row *Node) *Node

AddRow adds a 1xN row node to every row of n (bias broadcast).

func (*Node) Backward 

func (n *Node) Backward()

Backward runs reverse-mode differentiation from n, which should be a scalar (1x1) loss. Gradients accumulate into the Grad field of every contributing Param node.

func (*Node) MSELoss 

func (n *Node) MSELoss(target *Matrix) *Node

MSELoss returns the scalar mean squared error against a constant target.

func (*Node) MatMul 

func (n *Node) MatMul(o *Node) *Node

MatMul returns the matrix product n * o.

func (*Node) Mean 

func (n *Node) Mean() *Node

Mean reduces n to a 1x1 scalar averaging all elements.

func (*Node) MulElem 

func (n *Node) MulElem(o *Node) *Node

MulElem returns the element-wise (Hadamard) product. Shapes must match.

func (*Node) Named 

func (n *Node) Named(name string) *Node

Named sets a display name shown by ToDot and returns the node.

func (*Node) ReLU 

func (n *Node) ReLU() *Node

ReLU applies max(0, x) element-wise.

func (*Node) Scalar 

func (n *Node) Scalar() Float

Scalar returns the value of a 1x1 node.

func (*Node) Scale 

func (n *Node) Scale(s Float) *Node

Scale returns n multiplied by the scalar s.

func (*Node) Sigmoid 

func (n *Node) Sigmoid() *Node

Sigmoid applies 1/(1+e^-x) element-wise.

func (*Node) Softmax 

func (n *Node) Softmax() *Node

Softmax normalizes each row of n into a probability distribution.

func (*Node) SoftmaxCELoss 

func (n *Node) SoftmaxCELoss(target *Matrix) *Node

SoftmaxCELoss returns the scalar softmax cross-entropy against integer class labels (an Mx1 matrix of class indices), matching the SoftmaxCrossEntropy loss used by Sequential models.

func (*Node) Sub 

func (n *Node) Sub(o *Node) *Node

Sub returns the element-wise difference n - o. Shapes must match.

func (*Node) Sum 

func (n *Node) Sum() *Node

Sum reduces n to a 1x1 scalar by summing all elements.

func (*Node) T 

func (n *Node) T() *Node

T returns the transpose of n.

func (*Node) Tanh 

func (n *Node) Tanh() *Node

Tanh applies tanh(x) element-wise.

func (*Node) ToDot 

func (n *Node) ToDot() string

ToDot renders the computation graph rooted at n in Graphviz DOT format, in the spirit of Gorgonia's encoding/dot. Pipe it through the dot tool to get an image: 

go run ./_example/dot | dot -Tsvg > graph.svg

Leaves are drawn as boxes (Param blue, Input gray) and operations as rounded nodes; every node shows its shape. Use Named to label leaves.

type Optimizer 

type Optimizer interface {
	// Step applies one update to the given parameters using their gradients.
	// It is called once per parameterized layer.
	Step(idx int, weights, gradW *Matrix, bias, gradB []Float)
	// NewLayer registers a new parameterized layer and returns its index.
	NewLayer() int
	// Name returns a short identifier.
	Name() string
}

Optimizer updates a set of (weights, bias) parameter pairs using their gradients. One Optimizer instance is shared by the model; each parameterized layer gets its own state buffer inside the optimizer.

type Q4Matrix added in v0.0.2 

type Q4Matrix struct {
	Rows, Cols int
	Q          []uint8 // row quads x 2*Cols, padded for 32-byte loads
	Scale      []Float
	// ScaleMin, when non-nil, switches the per-(group, column)
	// dequantization from the symmetric offset-binary form
	// scale*(nibble-8) to the asymmetric scale*nibble - min — the form
	// GGUF's Q4_K sub-blocks carry — with Scale unused. Each entry packs
	// the pair as bfloat16 (PackScaleMin), halving what the kernels
	// stream per group next to two float32 tables. Nil keeps the
	// symmetric form.
	ScaleMin []uint32
	Group    int // input rows per scale; 0 means the default q4Group (64)
}

Q4Matrix is a weight matrix quantized to 4 bits with one scale per (64-row group, output column). Rows are stored in interleaved quads of two bytes per column, tiled 32 columns at a time: tile j/32 packs its row quads back to back (64 bytes apiece, see Index), so a kernel worker sweeping a tile range streams strictly sequential memory instead of striding across the full row width. Each byte holds two rows' nibbles (low nibble first), offset-binary (0..15 encodes -8..7), zero rows padding the final quad. The 256-bit kernel's nibble unpack turns those two bytes into four consecutive u8 lanes — exactly QMatrix's quad layout — so the same two-instruction multiply-add chain takes a column four rows deep, with activations re-centered to signed bytes and the nibble offset folded out through per-group activation sums.

func NewQ4Matrix added in v0.0.2 

func NewQ4Matrix(rows, cols, group int, minForm bool) *Q4Matrix

NewQ4Matrix allocates the layout for rows x cols with `group` input rows per scale (0 for the default 64); minForm picks the packed asymmetric scale/min table over the symmetric Scale. The caller fills Q via Index and the table it asked for.

func QuantizeMatrix4 added in v0.0.2 

func QuantizeMatrix4(m *Matrix) (*Q4Matrix, error)

QuantizeMatrix4 quantizes group-wise, symmetric with round-to-nearest. Columns split across CPUs for large matrices, like QuantizeMatrix.

func (*Q4Matrix) Index added in v0.0.2 

func (q *Q4Matrix) Index(i, j int) int

Index returns the position in Q of the byte carrying column j of rows i and i+1 (the low and the high nibble; shift by 4*(i%2)).

func (*Q4Matrix) MatMul added in v0.0.2 

func (q *Q4Matrix) MatMul(x, out *Matrix) error

MatMul computes out = x @ Q for a batch of activation rows — the prompt-prefill shape, mirroring QMatrix.MatMul: rows quantize to the same 7-bit form MatVec uses and the kernel processes them in blocks of four against one streaming pass over the nibbles.

func (*Q4Matrix) MatVec added in v0.0.2 

func (q *Q4Matrix) MatVec(x, out []Float) error

MatVec computes out = x @ Q for a single activation row: len(x) must be Rows and len(out) Cols. The activation row quantizes once per call, with its per-group sums carrying the nibble offset correction.

func (*Q4Matrix) TableIndex added in v0.0.2 

func (q *Q4Matrix) TableIndex(g, j int) int

TableIndex returns the position in Scale or ScaleMin of group g, column j. Tables are tile-major like the nibbles — tile, then group, then the 32 columns — so a kernel worker's table walk is sequential.

type Q8GMatrix added in v0.0.2 

type Q8GMatrix struct {
	Rows, Cols int
	Q          []int8 // interleaved row quads, padded for 32-byte loads
	Scale      []Float
	ColSum64   []int32 // per (group, column): 64 * sum of the group's weights
	Group      int     // input rows per scale; 0 means the default q8Group (32)
}

Q8GMatrix is a weight matrix quantized to int8 with one scale per (32-row group, output column) — the granularity llama.cpp's Q8_0 blocks carry, so a GGUF checkpoint loads by transposing bytes, never touching float32 weights. Rows are stored in the same interleaved quads as QMatrix; the kernels are the QMatrix kernels with the group scale and activation-offset correction folded in at group boundaries, Q4Matrix style.

func NewQ8GMatrix added in v0.0.2 

func NewQ8GMatrix(rows, cols, group int) *Q8GMatrix

NewQ8GMatrix allocates the layout for rows x cols with `group` input rows per scale (0 for the default 32); the caller fills Q (quad layout), Scale, and ColSum64 — see the loaders in _example/qwen.

func (*Q8GMatrix) Index added in v0.0.2 

func (q *Q8GMatrix) Index(i, j int) int

Index returns the position in Q of row i, column j: 32-column tiles store their row quads back to back (128 bytes apiece), so a kernel worker streams sequential memory.

func (*Q8GMatrix) MatMul added in v0.0.2 

func (q *Q8GMatrix) MatMul(x, out *Matrix) error

MatMul computes out = x @ Q for a batch of activation rows, in blocks of eight rows per weight stream like QMatrix.MatMul.

func (*Q8GMatrix) MatVec added in v0.0.2 

func (q *Q8GMatrix) MatVec(x, out []Float) error

MatVec computes out = x @ Q for a single activation row: len(x) must be Rows and len(out) Cols.

func (*Q8GMatrix) TableIndex added in v0.0.2 

func (q *Q8GMatrix) TableIndex(g, j int) int

TableIndex returns the position in Scale and ColSum64 of group g, column j; tables are tile-major like the weights.

type QMatrix added in v0.0.2 

type QMatrix struct {
	Rows, Cols int
	Q          []int8 // interleaved row quads, padded for 32-byte loads
	Scale      []Float
	ColSum64   []int32
}

QMatrix is a weight matrix quantized to int8 with one scale per output column: W[i][j] ~= Float(q_ij) * Scale[j]. Rows are stored in interleaved quads — Q[(i/4)*4*Cols + 4*j + i%4] holds rows i..i+3 of column j in four consecutive bytes, zero rows padding the final quad — which is the operand layout of the 256-bit u8 x s8 pairwise multiply-add followed by the widening i16 pair-add: two instructions take a column four rows deep.

Activations quantize per call to 7 bits with a +64 offset (see quantizeActs): the unsigned operand of the multiply then stays within [0,127], so the i16 pair sums cannot saturate, and the offset folds out through ColSum64, the precomputed per-column weight sums times 64.

func QuantizeMatrix added in v0.0.2 

func QuantizeMatrix(m *Matrix) *QMatrix

QuantizeMatrix quantizes column-wise, symmetric around zero with round-to-nearest. Columns are independent, so large matrices split across CPUs — quantize-at-load of a whole checkpoint is bound by this.

func (*QMatrix) Index added in v0.0.2 

func (q *QMatrix) Index(i, j int) int

Index returns the position in Q of row i, column j: 32-column tiles store their row quads back to back (128 bytes apiece), so a kernel worker streams sequential memory.

func (*QMatrix) MatMul added in v0.0.2 

func (q *QMatrix) MatMul(x, out *Matrix) error

MatMul computes out = x @ Q for a batch of activation rows — the prompt-prefill shape. Each row quantizes to the same 7-bit form MatVec uses, and the kernel processes rows in blocks of eight against one streaming pass over the weights, so the weight traffic that dominates a single matvec amortizes across the batch.

func (*QMatrix) MatVec added in v0.0.2 

func (q *QMatrix) MatVec(x, out []Float) error

MatVec computes out = x @ Q for a single activation row: len(x) must be Rows and len(out) Cols. The activation row is quantized once per call; output columns split across CPUs.

type RNNCell 

type RNNCell struct {
	Wx *Node // in x hidden
	Wh *Node // hidden x hidden
	B  *Node // 1 x hidden
}

RNNCell is a simple (Elman) recurrent cell: h' = tanh(x*Wx + h*Wh + b).

func NewRNNCell 

func NewRNNCell(inSize, hidden int, rng *rand.Rand) *RNNCell

NewRNNCell returns a randomly initialized RNN cell.

func (*RNNCell) InitState 

func (c *RNNCell) InitState(batch int) *Node

InitState returns a zero hidden state for the given batch size.

func (*RNNCell) Params 

func (c *RNNCell) Params() []*Node

Params returns the cell's trainable parameters, for NewTrainer.

func (*RNNCell) Step 

func (c *RNNCell) Step(x, h *Node) *Node

Step consumes one time step and the previous hidden state, returning the next hidden state.

type ReLU 

type ReLU struct {
	// contains filtered or unexported fields
}

ReLU activation: f(x) = max(0, x).

func (*ReLU) Backward 

func (r *ReLU) Backward(gradOutput *Matrix) (*Matrix, error)

func (*ReLU) Forward 

func (r *ReLU) Forward(input *Matrix) (*Matrix, error)

func (*ReLU) Grads 

func (a *ReLU) Grads() (*Matrix, []Float)

func (*ReLU) Init 

func (r *ReLU) Init(inputCols int, _ *rand.Rand) (int, error)

func (*ReLU) Params 

func (a *ReLU) Params() (*Matrix, []Float)

func (*ReLU) SetParams 

func (a *ReLU) SetParams(*Matrix, []Float) error

type SGD 

type SGD struct {
	LR       Float
	Momentum Float
	// contains filtered or unexported fields
}

SGD is stochastic gradient descent with optional momentum.

func NewSGD 

func NewSGD(lr, momentum Float) *SGD

NewSGD returns an SGD optimizer. Momentum of 0 disables momentum.

func (*SGD) Name 

func (s *SGD) Name() string

Name returns "sgd".

func (*SGD) NewLayer 

func (s *SGD) NewLayer() int

NewLayer registers a layer and returns its index.

func (*SGD) Step 

func (s *SGD) Step(idx int, weights, gradW *Matrix, bias, gradB []Float)

Step updates parameters with standard (momentum) SGD.

func (*SGD) String 

func (s *SGD) String() string

String returns a human-readable description of an optimizer's config.

type SelfAttention 

type SelfAttention struct {
	Wq, Wk, Wv *Node // inSize x dModel
}

SelfAttention is a single-head self-attention block with learned query, key, and value projections. It operates on one sequence at a time: the input is a (seqLen x inSize) node.

func NewSelfAttention 

func NewSelfAttention(inSize, dModel int, rng *rand.Rand) *SelfAttention

NewSelfAttention returns a randomly initialized self-attention block.

func (*SelfAttention) Forward 

func (a *SelfAttention) Forward(x *Node) *Node

Forward applies self-attention to a (seqLen x inSize) sequence, returning a (seqLen x dModel) sequence.

func (*SelfAttention) Params 

func (a *SelfAttention) Params() []*Node

Params returns the block's trainable parameters, for NewTrainer.

type Sequential 

type Sequential struct {
	// contains filtered or unexported fields
}

Sequential stacks layers and runs forward/backward passes.

func NewSequential 

func NewSequential() *Sequential

NewSequential returns an empty Sequential model. optimizer and loss are configured via Compile.

func (*Sequential) Add 

func (s *Sequential) Add(layer Layer) *Sequential

Add appends a layer to the network. Layers are added in forward order.

func (*Sequential) Compile 

func (s *Sequential) Compile(inputCols int, loss Loss, optimizer Optimizer) error

Compile wires the loss and optimizer and initializes all parameters. inputCols is the number of features in a single input row.

func (*Sequential) Fit 

func (s *Sequential) Fit(input, target *Matrix, epochs int) error

Fit trains the model for the given number of epochs over the dataset. If epochs > 1 the full dataset is reused each epoch (full-batch by default; callers can pass minibatches to FitStep directly for finer control).

func (*Sequential) FitStep 

func (s *Sequential) FitStep(input, target *Matrix) (Float, error)

FitStep performs one forward + loss + backward + update pass for a batch and returns the average loss for that batch.

func (*Sequential) Layers 

func (s *Sequential) Layers() []Layer

Layers returns the layers in forward order, for tools that walk the model structure (e.g. format exporters).

func (*Sequential) Load 

func (s *Sequential) Load(r io.Reader) error

Load restores parameters saved by Save into a model compiled with the same architecture.

func (*Sequential) LoadFile 

func (s *Sequential) LoadFile(path string) error

LoadFile restores parameters from a file written by SaveFile.

func (*Sequential) Predict 

func (s *Sequential) Predict(input *Matrix) (*Matrix, error)

Predict runs a forward pass with no gradient tracking.

func (*Sequential) Save 

func (s *Sequential) Save(w io.Writer) error

Save writes the model's parameters as JSON. The architecture itself is not stored: Load must be called on a model built and compiled with the same layers.

func (*Sequential) SaveFile 

func (s *Sequential) SaveFile(path string) error

SaveFile writes the model's parameters to a JSON file.

type Sigmoid 

type Sigmoid struct {
	// contains filtered or unexported fields
}

Sigmoid activation: f(x) = 1 / (1 + e^-x).

func (*Sigmoid) Backward 

func (s *Sigmoid) Backward(gradOutput *Matrix) (*Matrix, error)

func (*Sigmoid) Forward 

func (s *Sigmoid) Forward(input *Matrix) (*Matrix, error)

func (*Sigmoid) Grads 

func (a *Sigmoid) Grads() (*Matrix, []Float)

func (*Sigmoid) Init 

func (s *Sigmoid) Init(inputCols int, _ *rand.Rand) (int, error)

func (*Sigmoid) Params 

func (a *Sigmoid) Params() (*Matrix, []Float)

func (*Sigmoid) SetParams 

func (a *Sigmoid) SetParams(*Matrix, []Float) error

type Softmax 

type Softmax struct {
	// contains filtered or unexported fields
}

Softmax normalizes each row into a probability distribution. Unlike the element-wise activations its backward pass couples all columns of a row. Note that SoftmaxCrossEntropy already applies softmax internally; use this layer only when the model output itself must be probabilities (e.g. with a custom loss).

func (*Softmax) Backward 

func (s *Softmax) Backward(gradOutput *Matrix) (*Matrix, error)

func (*Softmax) Forward 

func (s *Softmax) Forward(input *Matrix) (*Matrix, error)

func (*Softmax) Grads 

func (a *Softmax) Grads() (*Matrix, []Float)

func (*Softmax) Init 

func (s *Softmax) Init(inputCols int, _ *rand.Rand) (int, error)

func (*Softmax) Params 

func (a *Softmax) Params() (*Matrix, []Float)

func (*Softmax) SetParams 

func (a *Softmax) SetParams(*Matrix, []Float) error

type SoftmaxCrossEntropy 

type SoftmaxCrossEntropy struct{}

SoftmaxCrossEntropy combines softmax + cross-entropy with integer class labels. Targets must be an Mx1 matrix whose entries are class indices.

func (SoftmaxCrossEntropy) Loss 

func (SoftmaxCrossEntropy) Loss(pred, target *Matrix) (Float, *Matrix, error)

Loss returns the average negative log-likelihood of the target classes and the combined softmax-cross-entropy gradient (pred - onehot) / batch.

func (SoftmaxCrossEntropy) LossInto 

func (SoftmaxCrossEntropy) LossInto(pred, target, grad *Matrix) (Float, error)

LossInto writes the softmax-cross-entropy gradient into grad.

func (SoftmaxCrossEntropy) Name 

func (SoftmaxCrossEntropy) Name() string

Name returns "softmax_ce".

type Tanh 

type Tanh struct {
	// contains filtered or unexported fields
}

Tanh activation: f(x) = tanh(x).

func (*Tanh) Backward 

func (t *Tanh) Backward(gradOutput *Matrix) (*Matrix, error)

func (*Tanh) Forward 

func (t *Tanh) Forward(input *Matrix) (*Matrix, error)

func (*Tanh) Grads 

func (a *Tanh) Grads() (*Matrix, []Float)

func (*Tanh) Init 

func (t *Tanh) Init(inputCols int, _ *rand.Rand) (int, error)

func (*Tanh) Params 

func (a *Tanh) Params() (*Matrix, []Float)

func (*Tanh) SetParams 

func (a *Tanh) SetParams(*Matrix, []Float) error

type Tensor added in v0.0.2 

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 added in v0.0.2 

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 added in v0.0.2 

func NewTensor(shape ...int) *Tensor

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

func NewTensorFromSlice added in v0.0.2 

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

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

func (*Tensor) Add added in v0.0.2 

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

Add returns t + o element-wise with broadcasting.

func (*Tensor) At added in v0.0.2 

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

At returns the element at the given multi-index.

func (*Tensor) Div added in v0.0.2 

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 added in v0.0.2 

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

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

func (*Tensor) Mul added in v0.0.2 

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

Mul returns t * o element-wise with broadcasting.

func (*Tensor) Reshape added in v0.0.2 

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 added in v0.0.2 

func (t *Tensor) Scale(s Float)

Scale multiplies every element by s, in place.

func (*Tensor) Set added in v0.0.2 

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

Set sets the element at the given multi-index.

func (*Tensor) Size added in v0.0.2 

func (t *Tensor) Size() int

Size returns the total number of elements.

func (*Tensor) Sub added in v0.0.2 

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

Sub returns t - o element-wise with broadcasting.

func (*Tensor) Transpose added in v0.0.2 

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 added in v0.0.2 

func (t *Tensor) Validate() error

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

type Trainer 

type Trainer struct {
	// contains filtered or unexported fields
}

Trainer owns the optimizer bookkeeping for a set of autograd parameters, so a training step is just building the loss graph and calling Step. 

trainer := tensai.NewTrainer(tensai.NewAdam(0.05), w1, b1, w2, b2)
for step := 0; step < 2000; step++ {
	loss := forward(x).MSELoss(y)
	trainer.Step(loss)
}

func NewTrainer 

func NewTrainer(opt Optimizer, params ...*Node) *Trainer

NewTrainer registers the parameters with the optimizer and returns a Trainer that updates them.

func (*Trainer) Step 

func (t *Trainer) Step(loss *Node) Float

Step runs backward from the scalar loss, applies one optimizer update to every parameter, clears the gradients, and returns the loss value.

Source Files

View all Source files

Directories

Path Synopsis
_example
charrnn command
Command charrnn trains a character-level LSTM on a small embedded corpus and generates text from it.
 Command charrnn trains a character-level LSTM on a small embedded corpus and generates text from it.
dataset command
Command dataset walks through the Dataset workflow end to end: build a dataset, shuffle it, split off a test set, standardize using training statistics only, train with mini-batches, and evaluate on the held-out split.
 Command dataset walks through the Dataset workflow end to end: build a dataset, shuffle it, split off a test set, standardize using training statistics only, train with mini-batches, and evaluate on the held-out split.
dot command
Command dot prints the computation graph of z = x + y in Graphviz DOT format — tensai's equivalent of Gorgonia's encoding/dot example.
 Command dot prints the computation graph of z = x + y in Graphviz DOT format — tensai's equivalent of Gorgonia's encoding/dot example.
fizzbuzz command
gpt2 command
Command gpt2 runs the real, published GPT-2 small (124M) checkpoint in pure Go: the weights load through tensai's encoding/safetensors reader, the text goes through tensai's tokenizer package (byte-level BPE from tokenizer.json), and every matvec in the transformer runs on tensai's Dot kernel — build with GOEXPERIMENT=simd for the AVX2 version.
 Command gpt2 runs the real, published GPT-2 small (124M) checkpoint in pure Go: the weights load through tensai's encoding/safetensors reader, the text goes through tensai's tokenizer package (byte-level BPE from tokenizer.json), and every matvec in the transformer runs on tensai's Dot kernel — build with GOEXPERIMENT=simd for the AVX2 version.
helloworld command
Command helloworld is the smallest possible tensai program: build a computation graph that adds two values, evaluate it, and differentiate it — tensai's equivalent of Gorgonia's hello world.
 Command helloworld is the smallest possible tensai program: build a computation graph that adds two values, evaluate it, and differentiate it — tensai's equivalent of Gorgonia's hello world.
iris command
mnist command
plasma command
Command plasma renders a demoscene-style plasma effect in the terminal — except the plasma function is a neural network.
 Command plasma renders a demoscene-style plasma effect in the terminal — except the plasma function is a neural network.
qwen command
Command qwen runs the published Qwen2.5-0.5B-Instruct checkpoint in pure Go: BF16 weights load through encoding/safetensors, text goes through the tokenizer package, and the Qwen2 architecture — RMSNorm, rotary embeddings, grouped-query attention, SwiGLU — decodes with a KV cache.
 Command qwen runs the published Qwen2.5-0.5B-Instruct checkpoint in pure Go: BF16 weights load through encoding/safetensors, text goes through the tokenizer package, and the Qwen2 architecture — RMSNorm, rotary embeddings, grouped-query attention, SwiGLU — decodes with a KV cache.
spiral command
tensor command
Command tensor tours the n-dimensional Tensor API: NumPy-style broadcasting, batched matrix multiplication with a shared weight, and scaled dot-product attention over a whole batch in three lines.
 Command tensor tours the n-dimensional Tensor API: NumPy-style broadcasting, batched matrix multiplication with a shared weight, and scaled dot-product attention over a whole batch in three lines.
wgpu command
Command wgpu exercises the experimental WebGPU backend: it reports the adapter wgpu-native picked, checks a GPU MatMul against the CPU one, and times both.
 Command wgpu exercises the experimental WebGPU backend: it reports the adapter wgpu-native picked, checks a GPU MatMul against the CPU one, and times both.
xor command
cmd
tensai command
Command tensai runs GGUF and safetensors language models on tensai's pure-Go kernels.
 Command tensai runs GGUF and safetensors language models on tensai's pure-Go kernels.
encoding
gguf
Package gguf reads the GGUF model format (llama.cpp's container: https://github.com/ggml-org/ggml/blob/master/docs/gguf.md) — typed metadata key/values followed by an aligned blob of tensors — with no dependencies beyond the standard library.
 Package gguf reads the GGUF model format (llama.cpp's container: https://github.com/ggml-org/ggml/blob/master/docs/gguf.md) — typed metadata key/values followed by an aligned blob of tensors — with no dependencies beyond the standard library.
onnx
Package onnx marshals trained tensai Sequential models into the ONNX format (opset 13, FP32, batch size 1), with the protobuf writer implemented in-tree — no dependencies.
 Package onnx marshals trained tensai Sequential models into the ONNX format (opset 13, FP32, batch size 1), with the protobuf writer implemented in-tree — no dependencies.
safetensors
Package safetensors reads and writes the safetensors checkpoint format (https://github.com/huggingface/safetensors) — the plain "8-byte header length, JSON header, raw little-endian buffer" layout most published model weights ship in — with no dependencies beyond the standard library.
 Package safetensors reads and writes the safetensors checkpoint format (https://github.com/huggingface/safetensors) — the plain "8-byte header length, JSON header, raw little-endian buffer" layout most published model weights ship in — with no dependencies beyond the standard library.
tflite
Package tflite marshals trained tensai Sequential models into the TensorFlow Lite FlatBuffers format (FP32, batch size 1), so they can run on the TFLite / LiteRT runtimes — including from Go via github.com/mattn/go-tflite (alias one of the packages when importing both, e.g.
 Package tflite marshals trained tensai Sequential models into the TensorFlow Lite FlatBuffers format (FP32, batch size 1), so they can run on the TFLite / LiteRT runtimes — including from Go via github.com/mattn/go-tflite (alias one of the packages when importing both, e.g.
internal
llm
Package llm wires tensai's kernels into a runnable language model: checkpoint download and loading, chat templates, sampling, generation (plain and speculative), the GPU decode path, and the OpenAI-compatible server.
 Package llm wires tensai's kernels into a runnable language model: checkpoint download and loading, chat templates, sampling, generation (plain and speculative), the GPU decode path, and the OpenAI-compatible server.
mmapfile
Package mmapfile memory-maps files read-only, so checkpoint readers can slice tensor bytes straight out of the page cache instead of copying them through read buffers.
 Package mmapfile memory-maps files read-only, so checkpoint readers can slice tensor bytes straight out of the page cache instead of copying them through read buffers.
Package tokenizer loads Hugging Face tokenizer.json files and implements the byte-level BPE family they describe — the tokenizers of GPT-2, Llama 3, Qwen, and most other published byte-level models — with no dependencies.
 Package tokenizer loads Hugging Face tokenizer.json files and implements the byte-level BPE family they describe — the tokenizers of GPT-2, Llama 3, Qwen, and most other published byte-level models — with no dependencies.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.