GoMLX, an Accelerated ML and Math Framework

📖 About GoMLX

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

[!Tip]

• See our 🎓 tutorial 🎓
• See Eli Bendersky’s blog post “GoMLX: ML in Go without Python”, (a bit outdated, but still useful)
• A guided example for Kaggle Dogs Vs Cats.
• A simple GoMLX slide deck with small sample code.

🔮 Upcoming Features & Plans 🔮

Large API and package re-organization coming soon in v0.28 release:

• backends moved to github.com/gomlx/compute repository!
  • Packages backends, dtypes, shapes and distributed moved to github.com/gomlx/compute.
  • The “go” backend is now implemented in github.com/gomlx/compute/gobackend (it’s been greatly improved as well).
  • The “xla” backend is now implemented in github.com/gomlx/go-xla/compute/xla.
• Removed pkg/ prefix from the the top level packages (no need with internal/ special handling).
• The context.Context variable container redesigned (and simplified) into model.Store (the container) the model.Scope (a pointer to a Store with the current scope).

See more details in the our CHANGELOG, including how to use cmd/convert_v0.28, a tool that facilitates the changes needed to move to the new API – details in the CHANGELOG.

Upcoming features for the upcoming v0.28.0 release:

• Dynamic shapes support for GoMLX: alpha/experimental, only for the Go backend (XLA only supports static shapes). But it will allow more efficient and padding-free (sometimes) models with Go. It also opens up better support for ONNXRuntime backend (a goal for after the v0.28 release).
• Improvements and some SIMD support for the Go backend (now in repo github.com/gomlx/compute/gobackend).

🗺️ Overview

GoMLX is a full-featured ML framework, supporting various well-known ML components
from the bottom to the top of the stack. But it is still only a slice of what a major ML library/framework should provide (like TensorFlow, Jax, or PyTorch).

Examples developed using GoMLX

• 🚀 NEW 🚀 KaLM-Gema3 12B parameters: Tecent’s top-ranked sentence encoder for RAGs, using go-huggingface to load the model and tokenizer, and GoMLX to execute it.
• 🚀 NEW 🚀 Gemma 3 270M: Demonstrates ONNX-converted text generation (LLM) using the onnx-community/gemma-3-270m-it-ONNX model with GoMLX. It uses the gomlx/onnx-gomlx package to convert the model, and gomlx/go-huggingface to download the model and run the * 🚀 NEW 🚀 GPT-2: Demonstrates text generation using the the new (experimental) transformer and generator packages. tokenizer.
• 🚀 NEW 🚀 BERT-base-NER: A BERT-base model fine-tuned for Named Entity Recognition. It’s also a ONNX-converted model from dslim/bert-base-NER model from HuggingFace.

• 🚀 NEW 🚀 MixedBread Reranker v1: A cross-encoder reranking example, see HuggingFace MixedBread Reranker v1 page. It uses the gomlx/onnx-gomlx package to convert the model, and gomlx/go-huggingface to download the model and run the tokenizer.

• Adult/Census model;
• How do KANs learn ?;
• Cifar-10 demo;
• MNIST demo (library and command-line only)
• Dogs & Cats classifier demo;
• IMDB Movie Review demo;
• Diffusion model for Oxford Flowers 102 dataset (generates random flowers);
  • Flow Matching Study Notebook based on Meta’s “Flow Matching Guide and Code”.
• GNN model for OGBN-MAG (experimental).
• Last, a trivial synthetic linear model, for those curious to see a barebones simple model.
• Neural Style Transfer 10-year Celebration: see a demo written using GoMLX of the original paper.
• Triplet Losses: various negative sampling strategies as well as various distance metrics.
• AlphaZero AI for the game of Hive: it uses a trivial GNN to evaluate positions on the board. It includes a WASM demo (runs GoMLX in the browser!) and a command-line UI to test your skills!

Backends

This page is excerpted from the full README. For complete documentation, browse the sections in the sidebar.