Overview

GoMLX is an easy-to-use set of Machine Learning and generic math libraries and tools. It can be seen as a PyTorch/Jax/TensorFlow for Go.

It can be used to train, fine-tune, modify, and combine machine learning models (it reads models from HuggingFace, with a growing list of model support). It provides all the tools to make that work easy: from a complete set of differentiable operators, all the way to UI tools to plot metrics while training in a notebook.

It defines a common “backend” API to run models. It includes a pure Go backend that is portable and also runs in WASM (in a browser), see demo created with GoMLX.

The “xla” backend is an optimized engine based on OpenXLA that uses just-in-time compilation to CPU, GPUs (Nvidia, and likely AMD ROCm, Intel, Macs) and Google’s TPUs. It also supports modern distributed execution (new, still being actively improved) for multi-TPU or multi-GPU using XLA Shardy, an evolution of the GSPMD distribution). It’s the same engine that powers Google’s Jax, TensorFlow and Pytorch/XLA, and it has the same speed in many cases (*).

More recently, it added the “onnx” backend, which uses ONNX Runtime to run GoMLX computations. It can also save models to .onnx file format.

Highlights

• 🚀 NEW 🚀: Gradient checkpointing: trade-off memory usage for recomputation when training large models, with a very simple API.

• 🚀 NEW 🚀: Save models to a .onnx file, that can be used with ONNX Runtime. See example in UCI-Adult demo.

• HuggingFace Go compatibility with go-huggingface:

  • Download files from models/datasets sharing the same cache framework as the python version.
  • Tokenizers for various classes in pure Go, downloaded directly from HuggingFace.
  • Datasets iterators (using Apache’s Parquet format), to experiment with standard datasets.
  • Model parameters reading from GGUF or safetensors format.
  • Model conversion to GoMLX (some models at least) with a compatible transformer library. Includes support to sentence embedding (equivalent to sentence_transformer Python library).

• Convert ONNX models to GoMLX with onnx-gomlx: both as an alternative for onnxruntime (leveraging XLA), but also to further fine-tune models.

• Docker “gomlx_jupyterlab” with integrated JupyterLab and GoNB (a Go kernel for Jupyter notebooks)

• Autodiff: automatic differentiation—only gradients for now, no jacobian.

• Store and Scope: simple variable management for ML models.

• ML layers library with the most popular machine learning “layers”: FFN layers,
  various activation functions, layer and batch normalization, DyT (Dynamic Tanh normalizer, https://arxiv.org/abs/2503.10622), convolutions, pooling, dropout, Multi-Head-Attention (for transformer layers), LSTM, KAN (B-Splines, GR-KAN/KAT networks, Discrete-KAN, PiecewiseLinear KAN), PiecewiseLinear (for calibration and normalization), various regularizations, FFT (reverse/differentiable), learnable rational functions (both for activations and GR-KAN/KAT networks), VNN (Vector Neural Networks) for SO(3)-Equivariant/Invariant layers, etc.

• Training library, with some pretty-printing:

  • Plots for Jupyter notebook, using GoNB, a Go Kernel.
  • Various debugging tools: collecting values for particular nodes for plotting, simply logging the value of nodes during training, stack-trace of the code where nodes are created.

• gomlx_checkpoints, the command line tool to inspect checkpoint of train(-ing) models, generate plots with loss and arbitrary evaluation metrics using Vizb (a separate CLI tool – --plot is the only flag that needs it installed; everything else works without it). See cmd/gomlx_checkpoints/README.md for usage details. It also allows plotting different models together, to compare their evolution, and -loop gives you a live, auto-refreshing view while a model is still training.

• Various optimizers: SGD, Adam (AdamW and Adamax).

• Various losses and metrics.

• Read Numpy arrays into GoMLX tensors – see package github.com/gomlx/gomlx/core/tensors/numpy.

• Distributed Execution (experimental) across multiple GPUs or TPUs with little hints from the user. One only needs to configure a distributed dataset, and the trainer picks up from there. See code change in UCI-Adult demo. Experimental, please report any issues and help us improve it.