Running ML in Go Without Python
A previous article here walked through running machine learning inference in Go by calling out to a Python sidecar process. But that approach still depends on Python somewhere in the stack. GoMLX, a relatively new package for ML in Go, aims to eliminate that dependency entirely by sitting directly on the XLA compiler and PJRT runtime — the same infrastructure used by TensorFlow and JAX.
Where GoMLX fits in the ML stack
Most ML models are written in Python using frameworks like TensorFlow, JAX, or PyTorch. These frameworks handle two key concerns: providing an expressive API for model architecture (including automatic differentiation for training), and implementing fast computation primitives across CPUs, GPUs, and TPUs.
Between those layers sits the OpenXLA system, a standardized approach to model definition and hardware execution:
- The top layer is composed of frameworks that translate model definitions into StableHLO, a common interchange format for high-level operations.
- The bottom layer is the hardware that actually executes the models.
- In the middle sits OpenXLA's two core components: the XLA compiler, which turns HLO into machine code, and PJRT, the runtime that manages devices, data movement, and task execution.
That middle and bottom layer is implemented in C and C++, not Python. Python only appears in the frameworks at the very top. GoMLX exploits this by placing itself right where a Python framework would normally be — it provides Go-level primitives for building and training models, then lets XLA and PJRT handle the rest.
Building a CNN for CIFAR-10 in Go
To see how GoMLX works in practice, consider a convolutional neural network for the CIFAR-10 dataset. The model is defined as a computational graph using GoMLX's builder API. Operations are threaded explicitly through builder calls, making the code more verbose than Python's Keras-style interfaces, but a higher-level library could easily be layered on top.
The full graph definition is verbose but familiar. The training loop makes use of GoMLX's automatic differentiation support to compute gradients, and the model trains with similar results to the equivalent TF+Keras implementation from the previous article.
Inference is straightforward as well. A classifier function takes an image.Image, runs it through the network, and returns an index into the list of CIFAR-10 labels:
func main() {
flagCheckpoint := flag.String("checkpoint", "", "Directory to load checkpoint from")
flag.Parse()
mlxctx := mlxcontext.New()
backend := backends.New()
_, err := checkpoints.Load(mlxctx).Dir(*flagCheckpoint).Done()
if err != nil {
panic(err)
}
mlxctx = mlxctx.Reuse() // helps sanity check the loaded context
exec := mlxcontext.NewExec(backend, mlxctx.In("model"), func(mlxctx *mlxcontext.Context, image *graph.Node) *graph.Node {
// Convert our image to a tensor with batch dimension of size 1, and pass
// it to the C10ConvModel graph.
image = graph.ExpandAxes(image, 0) // Create a batch dimension of size 1.
logits := cnnmodel.C10ConvModel(mlxctx, nil, []*graph.Node{image})[0]
// Take the class with highest logit value, then remove the batch dimension.
choice := graph.ArgMax(logits, -1, dtypes.Int32)
return graph.Reshape(choice)
})
// classify takes a 32x32 image and returns a Cifar-10 classification according
// to the models. Use C10Labels to convert the returned class to a string
// name. The returned class is from 0 to 9.
classify := func(img image.Image) int32 {
input := images.ToTensor(dtypes.Float32).Single(img)
outputs := exec.Call(input)
classID := tensors.ToScalar[int32](outputs[0])
return classID
}
// ...
Running Gemma2 with GoMLX
A more demanding test of GoMLX's capabilities is its full implementation of Gemma2 inference. The model lives in the transformers package of the gomlx/gemma repository and follows a standard transformer architecture that will look instantly recognizable to anyone who has seen such models implemented in Python or elsewhere.
With weights downloaded from HuggingFace or Kaggle, the model can be loaded directly in Go. No Python process is needed to run the inference:
var (
flagDataDir = flag.String("data", "", "dir with converted weights")
flagVocabFile = flag.String("vocab", "", "tokenizer vocabulary file")
)
func main() {
flag.Parse()
ctx := context.New()
// Load model weights from the checkpoint downloaded from Kaggle.
err := kaggle.ReadConvertedWeights(ctx, *flagDataDir)
if err != nil {
log.Fatal(err)
}
// Load tokenizer vocabulary.
vocab, err := sentencepiece.NewFromPath(*flagVocabFile)
if err != nil {
log.Fatal(err)
}
// Create a Gemma sampler and start sampling tokens.
sampler, err := samplers.New(backends.New(), ctx, vocab, 256)
if err != nil {
log.Fatalf("%+v", err)
}
start := time.Now()
output, err := sampler.Sample([]string{
"Are bees and wasps similar?",
})
if err != nil {
log.Fatalf("%+v", err)
}
fmt.Printf("\tElapsed time: %s\n", time.Since(start))
fmt.Printf("Generated text:\n%s\n", strings.Join(output, "\n\n"))
}
The fact that GoMLX can run a production-grade open LLM entirely in Go is evidence that the package has reached a meaningful level of maturity in terms of its core computation and runtime integration.
A pragmatic division of labor
GoMLX reuses the hardest parts of the ML stack — the compiler, device runtime, and kernel implementations — and replaces only the model-building layer with a Go-native one. This avoids re-implementing the enormous body of optimized low-level libraries developed by Google, NVIDIA, Intel, and others.
Since GoMLX is a young project, it may not yet be suitable for every production scenario. But the approach is sound, and the ability to keep ML entirely within a Go process is a valuable step forward for teams that would rather avoid Python as a sidecar dependency.
