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 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 runs almost everywhere Go runs, using a pure Go backend. It runs even in the browser with WASM (see demo created with GoMLX). Likely, it will work in embedded devices as well (see Tamago).

It also supports a very optimized backend 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. Use this backend to train large models or with large datasets.

Highlights

• 🚀 NEW 🚀: Gradient checkpointing: trade-off memory usage for recomputation when training large models, with a very simple API.
• 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, 🚀 NEW 🚀 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 Plotly. See example of training session, with the effects of a learning rate change during the training. It also allows plotting different models together, to compare their evolution.
• 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, pls report any issues and help us improve it.