GoMLX: A Go Machine Learning Framework With Three Interchangeable Backends
GoMLX: An Accelerated Machine Learning Framework For Go
At a glance
- What is it?
- GoMLX puts a PyTorch-style graph API in Go and lets the same computation run on a pure Go interpreter, on OpenXLA, or on ONNX Runtime. The backend abstraction is the interesting part; the maturity of each backend is not equal.
- Who is it for?
- GoMLX is a reasonable choice if you already write Go and want training, fine-tuning or HuggingFace model loading inside the same binary as the rest of your service, and if you accept that the xla and onnx backends pull in C libraries while the go backend keeps you pure. It is not the right tool if you need a mature distributed training stack today, or if your team's existing model code is in Python and nobody wants a second implementation.
- Can I use it commercially?
- Yes. Apache-2.0 is a permissive licence: you can use, modify and sell software built on it, as long as you keep its copyright and licence notices.
- Is it still maintained?
- Yes. The repository last received commits 3 days ago.
- What is it written in?
- Mainly Go, according to GitHub's language statistics.
Answers come from the project's GitHub data, last synced on September 15, 2026, and from our analysis. They are not legal advice.
DEEP OPEN-SOURCE ANALYSIS
The gap GoMLX is trying to fill in the Go ecosystem
Go has plenty of HTTP servers and not much in the way of a full training framework. The options that exist tend to be either narrow (a single model type, or inference only) or bindings around a C library. GoMLX positions itself as "a PyTorch/Jax/TensorFlow for Go", which is a large claim, and the README is candid that it is "still only a slice of what a major ML library/framework should provide".
The audience is Go developers who want to define, train, fine-tune or combine models without leaving the language. The README lists training, fine-tuning, modifying and combining models, plus reading models from HuggingFace with "a growing list of model support". The examples named in the README are the usual entry points: UCI Adult/Census, Cifar-10, MNIST, Dogs vs Cats, IMDB reviews, a KAN visualization notebook, and a diffusion model for Oxford Flowers 102. That mix says the target is someone doing tabular work, small vision models and text classification, not someone pretraining a foundation model.
The secondary audience is people who want to ship a model inside a Go binary. The pure Go backend runs in WebAssembly, and the README links a browser demo built with GoMLX, so putting a model in a web page without a Python service is an explicit use case.
One graph API, three backends with very different properties
The architectural decision that shapes everything else is the backend API. GoMLX defines a common interface for running models, and ships three implementations with distinct runtime characteristics.
The go backend is a pure Go interpreter, portable and able to run in WASM in a browser. The README describes recent work on it: the execution engine was updated "to minimize synchronization costs", buffer management improved, and SIMD support broadened using the portable simd package in go1.27. Matrix multiplication is the weak spot. The README states that only AVX2 and AVX512 were optimized, that the microkernel was written in assembly because of pending issues in Go's simd work (golang/go issues 80829 and 78753), and that there is no Neon path for Apple hardware yet. The project asks for donations of Apple hardware to add one. So the performance story on the pure Go path is: good on x86 with AVX2 or AVX512, unoptimized elsewhere.
The xla backend is built on OpenXLA and uses just-in-time compilation to CPU, Nvidia GPUs, likely AMD ROCm, Intel, Macs and Google TPUs. The README notes it is the same engine behind Jax, TensorFlow and PyTorch/XLA, and claims "the same speed in many cases" with an asterisk the README does not resolve in the excerpt available. Distributed execution across multiple TPUs or GPUs uses XLA Shardy, described as "new, still being actively improved". Treat that as a warning label, not a feature bullet.
The onnx backend runs computations through ONNX Runtime and supports onnx:cpu, onnx:cuda and onnx:rocm, plus onnx:wasm, webgpu and experimental webnn for WebAssembly. It can also save a model to .onnx, which the README suggests serving through other ONNX-compatible inference systems such as KnightAnalytics Hugot. That export path is the most concrete interoperability story in the README.
Getting a first model running
The README does not include an install block in the excerpt provided, so the exact go get line and the backend selection API are things to read from the documentation site at gomlx.github.io rather than guess. What the material does establish is the shape of the setup work.
Backend choice is a string, based on the backend names the README lists: onnx:cpu, onnx:cuda, onnx:rocm, onnx:wasm, webgpu and webnn for the ONNX path. The xla backend is referred to as xla, with xla:cpu appearing in the README's performance comparison against the go backend. The pure Go backend is referred to as go. How those strings are passed to a context or config is not shown in the material, so verify it against the current docs.
The practical install decision is which of these you need. The go backend keeps the build pure Go, which is what makes the WASM target possible. The xla and onnx backends both depend on external C libraries (OpenXLA and ONNX Runtime respectively), so they change your build and deployment story: cross-compilation, container base images and CI runners all have to account for them. If you only need inference on CPU in a Go service, the go backend avoids that entire class of problem, at the cost of the unoptimized matrix multiplication paths noted above.
There is also a Docker image, janpfeifer/gomlx_jupyterlab, linked from the README badges. That is the lowest-friction way to try the notebooks without assembling the XLA toolchain yourself.
The API is young and has already broken once
The release history is the clearest signal about stability. v0.28.0 is described as a "Large API update; Lots of new features", and v0.28.11 followed roughly seven weeks later with backends and dynamic shapes. The version number is still 0.x. The README's own framing is that GoMLX is a slice of a full framework, not a finished one.
For anyone writing production code against GoMLX, that means pinning a version and reading release notes before upgrading is not optional. A release titled "Large API update" at 0.28 is a statement that the maintainers reserve the right to reshape the surface. The project's stated philosophy is that code should be "simple to read and reason about" and that documentation is kept current, which is a reasonable counterweight, but it does not change the versioning reality.
The second limitation is scope. The README lists a specific set of example models and says HuggingFace support has "a growing list of model support". Growing means not complete. If your plan depends on a particular architecture being loadable, that is a thing to confirm before committing, not after.
The third is the pure Go backend's hardware coverage. AVX2 and AVX512 only, no Neon. On an Apple laptop or an ARM server, the go backend is the portable fallback rather than the fast path, and you would be reaching for xla or onnx instead, which reintroduces the C dependency.
Distributed training is the fourth. The README labels multi-GPU and multi-TPU support as new and actively being improved. Building a training cluster on it now means building on something the project itself describes as in motion.
How this differs from Gorgonia and from ONNX Runtime bindings
The obvious Go alternative is Gorgonia, which is also a graph-based automatic differentiation library in pure Go. The difference in approach is the compilation model. Gorgonia builds and evaluates its own graph in Go. GoMLX instead defines a backend interface and can hand the graph to OpenXLA for JIT compilation, which is what makes GPU and TPU targets reachable at all. If your goal is CPU-only training on x86, the two are closer in spirit; if your goal is GPU, GoMLX has a path that Gorgonia does not.
The other alternative is binding ONNX Runtime directly and skipping GoMLX entirely. That works for inference, and it is a smaller dependency. What you lose is the graph-building and training side: GoMLX gives you differentiable operators, optimizers, and the ability to train or fine-tune before exporting. The onnx backend inside GoMLX is essentially the same runtime, but reached through GoMLX's API so that the same model code can also run on go or xla. If you only ever run inference on a fixed .onnx file, the direct binding is the leaner choice. If you want to train in Go and then export, GoMLX is the one with both halves.
Worth noting: exporting to .onnx is itself an escape hatch. You can train in GoMLX and serve through a different ONNX-compatible runtime, which means you are not locked into GoMLX's inference performance for the serving path.
Licence and the cost of keeping up
GoMLX is Apache-2.0, which is permissive and includes an explicit patent grant. For most commercial use that is a low-friction choice, but the usual caveat applies: the xla and onnx backends link against OpenXLA and ONNX Runtime, which carry their own licences, and those are what you actually need to review if you redistribute a binary. Nothing here is legal advice; read the dependency licences yourself.
The maintenance cost has two components. The first is the 0.x API churn described above: upgrades require reading release notes, and v0.28.0 shows that a large surface change can land in a minor version. The second is the build matrix. If you ship the xla or onnx backend, you are shipping a C dependency, and your CI has to build it for every target platform you support. The go backend is the only one that keeps your build as plain Go.
The README also points at a Slack channel (gophers.slack.com, channel gomlx) and a sponsor link. For a project at this stage, the channel is likely the fastest way to find out whether a specific HuggingFace architecture is supported, which is more useful than reading a model list that the README itself calls growing.
Who should pick this up and what to check first
Adopt GoMLX if you are writing Go, your models are in the small-to-medium range (tabular, image classification, text classification, small diffusion), and you want training and inference in the same binary and the same language. The backend interface is the reason to choose it over a raw ONNX Runtime binding: you can start on the pure Go backend for portability and move to xla or onnx when you need the hardware, without rewriting the model.
Do not adopt it if you need distributed training that is already stable, if your target hardware is ARM or Apple and you intended to rely on the pure Go backend's speed, or if your existing pipeline is Python and the team has no appetite for a second implementation of the same model. The README's own description of Shardy-based distribution as new and actively improving should be taken at face value.
Before writing production code, verify four things. First, that the backend you intend to ship builds and runs on your target OS and architecture, since the README's test badges cover Linux amd64 and arm64, Darwin arm64 and Windows amd64, and your target may be none of those. Second, that the specific HuggingFace model you need is supported, by asking on the Slack channel rather than inferring from the examples. Third, how the current API differs from any tutorial you find, given that v0.28.0 was a large API update. Fourth, whether the .onnx export path produces a file your serving runtime accepts, if that is how you plan to deploy. The last one is testable in an afternoon and settles the deployment question before it becomes an architectural one.
Editorial conclusion
GoMLX is a reasonable choice if you already write Go and want training, fine-tuning or HuggingFace model loading inside the same binary as the rest of your service, and if you accept that the xla and onnx backends pull in C libraries while the go backend keeps you pure. It is not the right tool if you need a mature distributed training stack today, or if your team's existing model code is in Python and nobody wants a second implementation. Before adopting, verify three things on your own machine: that the backend you intend to ship builds and runs under your target OS and architecture, that the HuggingFace model you care about is on the supported list, and what the v0.28.0 API break means for any code you write against the current docs.
Community notes