GoMLX tackles a problem many Go developers have felt keenly: running modern machine learning workflows without leaving the Go ecosystem or falling back on Python bindings. It aims to be a full-featured ML framework akin to PyTorch or TensorFlow, but written in pure Go, with support for hardware acceleration, automatic differentiation, and popular model architectures. The ability to train and run large language models such as GPT-2 and Gemma 3 entirely in Go — no Python dependency in sight — is a noteworthy achievement that sets it apart.
What GoMLX does and its architecture
GoMLX is a machine learning framework written in Go that provides the key building blocks needed to build, train, and deploy neural networks. It offers automatic differentiation, a comprehensive library of ML layers (including transformers, CNNs, RNNs, KANs, and graph neural networks), and a training library that integrates with Jupyter notebooks through GoNB, a Go kernel for Jupyter.
Under the hood, GoMLX supports multiple backends for computation:
An OpenXLA backend with just-in-time (JIT) compilation optimized for CPUs, GPUs, and TPUs. This backend brings high performance and hardware acceleration to GoMLX, leveraging the same OpenXLA infrastructure that powers frameworks like TensorFlow and JAX.
A pure-Go portable backend that compiles to WebAssembly (WASM), enabling ML models to run directly in browsers without native dependencies.
A Darwin/Metal backend optimized for Apple hardware.
This multi-backend approach allows GoMLX to be flexible across hardware environments, from edge devices running in WASM to high-performance servers with GPUs and TPUs.
In addition to training models from scratch, GoMLX can load and run ONNX models, broadening interoperability with models trained in other frameworks. Demonstrated examples include large language models like Gemma 3 (270 million parameters), GPT-2, BERT-base for named entity recognition, and diffusion models.
The project philosophy embraces Go’s strengths: simplicity, transparency, clear error messages with stack traces, and zero Python dependencies. This makes it appealing for Go developers who want to stay in a single language environment for their ML workflows.
Why GoMLX stands out technically
GoMLX stands out for bringing an advanced ML framework experience to a language that isn’t traditionally associated with machine learning. Go’s ecosystem generally lacks mature ML frameworks comparable to Python’s PyTorch or TensorFlow, making GoMLX a rare and ambitious project.
The automatic differentiation engine is a core strength, enabling gradient-based optimization necessary for training neural networks. The library of ML layers is comprehensive, covering a wide range of model types, including state-of-the-art architectures like transformers and GNNs.
Supporting multiple computation backends is a tradeoff that adds complexity but provides flexibility. The OpenXLA backend enables just-in-time compilation and hardware acceleration, positioning GoMLX to compete with Python frameworks on performance when running on accelerators. The pure-Go backend that compiles to WASM is particularly interesting — it opens doors for running ML inference directly in browsers or other WASM-supporting environments without any native dependencies.
The code quality is surprisingly clean for a project of this scope in Go, reflecting careful design decisions to keep things idiomatic and maintainable. Error messages with stack traces align well with Go’s developer experience expectations.
Tradeoffs include the relative immaturity of the Go ML ecosystem compared to Python. While GoMLX covers many features, the breadth of tooling, pre-trained models, community examples, and integrations won’t match Python’s. Additionally, the user base for GoMLX might be smaller, so support and third-party extensions could be limited.
Overall, GoMLX balances production-ready capabilities with Go idioms and cross-platform support, making it a noteworthy option for Go-centric ML workflows.
Quick start
The GoMLX project provides a straightforward path to getting started, especially with the pre-built Docker image that bundles GoMLX, JupyterLab, the GoNB kernel, and optional NVIDIA CUDA runtime for GPU support.
Here are the exact commands as documented:
docker pull janpfeifer/gomlx_jupyterlab:latest
docker run -it --rm -p 8888:8888 -v "${PWD}":/home/jupyter/work janpfeifer/gomlx_jupyterlab:latest
For GPU support, add the flag --gpus all to the docker run command.
This setup launches JupyterLab with GoMLX ready to use. You can then open the tutorial located at Projects/gomlx/examples/tutorial inside the container.
If you want to use only the pure Go backend (no XLA or GPU dependencies), you can simply import the backend in your Go code:
import _ "github.com/gomlx/gomlx/backends/simplego"
This requires no additional installation.
GoMLX also auto-installs required XLA PJRT plugins for CPU, GPU, and TPUs on Linux and Mac by default, simplifying setup for hardware acceleration.
How to explore the project
If you want to explore beyond the quick start, the repository includes extensive examples for training and inference with various models. The documentation and tutorial notebooks provide hands-on guidance.
The integration with GoNB (a Go kernel for Jupyter) is an interesting developer experience feature, allowing interactive experimentation within notebooks familiar to data scientists.
The backend architecture and support for ONNX models mean you can port models trained elsewhere and run them in GoMLX, a useful feature for production environments standardizing on Go.
Verdict
GoMLX is relevant for Go developers who want to build and train ML models without switching languages or relying on Python bindings. Its support for advanced model architectures, automatic differentiation, and multiple backends including OpenXLA acceleration and WASM runtime make it stand out in the Go ecosystem.
That said, it’s still early days compared to Python frameworks. The ecosystem, tooling, and community support are not as mature. For production ML workloads where Python is acceptable, existing frameworks might offer more convenience and a richer ecosystem.
However, if your project prioritizes Go for simplicity, deployment consistency, or integration reasons, GoMLX is worth serious consideration. The Docker image with Jupyter and GoNB lowers the barrier to experimentation.
In summary, GoMLX solves a real problem — bringing PyTorch/JAX-level ML capabilities to Go — with a clean codebase, multiple hardware backends, and a pragmatic approach to dependencies. It’s a solid option for Go-centric ML development, especially when you want to avoid Python altogether.
Related Articles
- onnxmltools: a Python toolkit for converting ML models to ONNX format — onnxmltools is a Python library for converting machine learning models from various frameworks into the ONNX format, ena
- Ollama: a unified CLI and API platform for local large language models — Ollama simplifies running and managing open-source large language models locally with a unified CLI and REST API, suppor
- vllm-mlx: Efficient LLM serving on Apple Silicon with SSD-tiered KV cache and continuous batching — vllm-mlx is a Python inference server for Apple Silicon that supports OpenAI and Anthropic APIs, featuring SSD-tiered KV
- ZenML: a unified MLOps platform bridging classical ML and AI agent orchestration — ZenML offers an open-source Python SDK to orchestrate full ML and AI agent lifecycles, integrating popular tools and ena
- LlamaFactory: modular, extensible fine-tuning framework for large language models — LlamaFactory offers a modular Python framework for fine-tuning 100+ LLMs with diverse algorithms and optimizations, incl
→ GitHub Repo: gomlx/gomlx ⭐ 1,405 · Go