Model Formats¶
tensai speaks four model formats, all with in-tree encoders and decoders — the FlatBuffers writer, the protobuf encoder, and both checkpoint readers are implemented in the repository, so the default build stays dependency-free.
| Package | Direction | Format |
|---|---|---|
encoding/safetensors |
read + write | The checkpoint format published model weights usually ship in |
encoding/gguf |
read | llama.cpp's model container, including the quantized ladder |
encoding/tflite |
write | .tflite flatbuffers for the TFLite/LiteRT runtimes |
encoding/onnx |
write | ONNX (opset 13, FP32) |
safetensors¶
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
Save/SaveFile write F32 checkpoints that the reference implementation reads back bit-for-bit; interoperability is verified in both directions.
GGUF¶
encoding/gguf reads llama.cpp's container the same lazy way: Open parses the typed metadata (String/Int/Float/KV) and the tensor directory, and each Tensor call reads and dequantizes just that tensor. Supported types: 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, the nonlinear IQ4_NL, and gpt-oss's MXFP4 — so the whole ladder of checkpoints usually published for llama.cpp opens directly, verified block-exact against real llama.cpp conversions.
Names, Info, and Metadata inspect a checkpoint without loading it. Dimensions come back row-major like every other reader here.
RoPE row order
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.
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 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:
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
ONNX export¶
import tensaionnx "github.com/mattn/tensai/encoding/onnx"
err := tensaionnx.MarshalFile("model.onnx", model)
Same layer support as the TFLite export, 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).