tensai

Core design and tensor model

tensai is a small machine learning framework meant for learning and experiments. It implements forward passes, backpropagation, and optimization in pure Go, and the default build has no external dependencies. The optional wgpu build tag adds exactly one, the cgo free ebitengine/purego binding, used to reach a WebGPU backend. At the data level the library centers on a Matrix type plus an N-dimensional Tensor that generalizes matrices to any rank, with element wise add, subtract, multiply, and divide operations that follow NumPy style broadcasting. Batched matrix multiplication runs the per matrix products on the same kernel as the Dot operation and parallelizes across the batch, while a Transpose permutes axes and a Reshape supports negative one inference. Tensors also offer zero copy views to and from Matrix, so conversions between the two representations do not force a copy. This design keeps the library approachable for someone reading the source, because the math types are ordinary Go values rather than wrappers around a larger tensor runtime. The layout section of the README maps each file to a responsibility, from tensor.go for matrix and vector ops to the SIMD and generic matmul kernels. The file layout in the README makes the responsibilities explicit, so a reader can open tensor.go for the matrix and vector operations and dot_simd.go for the AVX2 kernel without guessing where each piece lives.

Acceleration and training loop

Performance comes from SIMD kernels and careful allocation. The AVX2 matmul and element wise kernels are written with Go's experimental simd package, still pure Go with no cgo and no assembly files, and they build on amd64 with GOEXPERIMENT=simd while every other build falls back to portable code automatically. Matmul, the activation functions, LayerNorm, and the Adam update are all eight lane vectorized. The training path is written to stay out of the garbage collector: layers reuse their forward and backward scratch buffers across steps, and a full multilayer perceptron step runs in about twenty nine allocations, while Predict always returns fresh results. The layer set covers Embedding, Dense, Conv2D, MaxPool2D, BatchNorm, LayerNorm, and Dropout, plus the usual activations. Optimizers include momentum SGD, Adam, and AdamW with decoupled weight decay, and loss functions cover mean squared error, softmax cross entropy, and binary cross entropy. A Sequential model ties layers together with Compile, Fit or FitStep, and Predict, while a micrograd style reverse mode autograd engine handles models that do not fit the sequential mold and can render its graph for Graphviz. This autograd path means tensai is not limited to stacked layers, and a developer can build recurrent or attention based models on the same reverse mode engine that the sequential models use underneath.

Backends and interoperability

Beyond training, tensai focuses on reading and writing real model formats, which is what makes the small framework usable with published weights. The WebGPU backend builds with the wgpu tag on Linux, macOS, and Windows and runs batched matmul as a WGSL compute shader through wgpu native, still without cgo because the shared library is opened at runtime. Quantization support builds int8 and int4 weight only twins, with int4 group wise and int8 as a full integer path, reaching memory bandwidth that the README quantifies on sixteen cores. For serialization and exchange, the encoding packages read and write TFLite flatbuffers, read Hugging Face safetensors checkpoints and llama.cpp GGUF containers, and marshal Sequential models to ONNX with a hand written protobuf encoder. The tokenizer package loads Hugging Face tokenizer.json files and implements byte level BPE families plus SentencePiece built from GGUF vocabularies. Optional example programs show the range, from an XOR learner and an MNIST classifier to running the published GPT-2 124M checkpoint and a Qwen chat model in pure Go. These examples double as tests of the import and export paths against reference implementations. The example programs from an XOR learner to a GPT-2 runner also act as integration tests, because they exercise the import and export paths against the formats that real models use in production.

Editorial conclusion

The framework is released under the MIT license and its most recent commit was on 2026-08-26.