GoMLX Project Status Update: Early But Usable Machine Learning for Go (August 2025)

Forum topic · ✨步子哥 · 2025-09-24

Summary

A Chinese forum post reviews the current state of GoMLX, a Go machine learning framework built on OpenXLA/PJRT. As of August 2025, core training and inference work end-to-end on CPU and CUDA, with ResNet-18/34 trainable on a single GPU, though ResNet-50 awaits SeparableConv2D. Optimizers SGD, Adam, and AdamW are available; LAMB and Shampoo remain on the roadmap. Versions v0.3 to v0.4 cut CUDA kernel call overhead by 30%, fixed gradient flow bugs, and aligned MNIST/CIFAR-10 convergence with PyTorch. A NaN/Inf stack-trace debugging feature is merged but unreleased, and distributed training (mirror-sync data parallelism) runs small 2×A100 experiments on a dev branch. There is no native HuggingFace weight loading yet; the community onnx-gomlx tool converts PyTorch models via ONNX. Development is led almost entirely by one maintainer, Jan Pfeifer, with roughly 6-8 commits per month and ~15% external contributors. The verdict: GoMLX suits research, teaching, and edge single-binary deployments, but Python remains the safer choice for HuggingFace ecosystems or large-scale distributed training.

GoMLX Project Status Update (August 2025)

As of August 2025, GoMLX remains in an "early but usable" stage: the core training-inference pipeline works, but there are clear gaps before production readiness.

1. Feature Completeness

Basic operators: Fully-connected layers, CNNs, LSTM, Multi-Head Attention, KAN, and FFT are merged into the main branch. ResNet-18/34 can be fully trained on a single GPU, but ResNet-50 is still blocked on missing SeparableConv2D, which remains on the official TODO list.
Optimizers: SGD, Adam, and AdamW are usable. LAMB, Shampoo, and others are still on the roadmap.
Autodiff and JIT compilation: Backed by OpenXLA/PJRT, both CPU and CUDA backends work. For small batches of 32–128 samples, training speed is roughly on par with PyTorch on CPU.
Pretrained models: No ability to directly load HuggingFace weights yet. The community provides the onnx-gomlx tool — convert PyTorch models to ONNX first, then import them into GoMLX for inference.

2. Recent Progress (2024 H2 – 2025 H1)

v0.3 → v0.4 reduced CUDA kernel call overhead by 30% and fixed gradient flow interruption issues. MNIST/CIFAR-10 convergence curves in the examples now match official PyTorch scripts.
The author's April 2025 blog mentions work on "NaN/Inf first-occurrence stack traces" to ease debugging of large models; the feature is merged into main but not yet released.
A distributed training branch (mirror-sync data parallelism) can run small experiments on 2×A100, but the code is still on the dev channel and the API may change.

3. Community and Ecosystem

The main repository averages 6–8 commits per month since September 2024. Most PRs come from the author himself; external contributors account for about 15%.
Tutorials: More than 10 ipynb notebooks (compatible with GoNB) can be run directly in Jupyter. A beginner can reproduce MNIST classification in about 30 minutes — a noticeably lower barrier to entry than Gorgonia.
Third-party adoption: Only small edge projects so far (IoT visual inspection, K8s inference pods). No company has publicly claimed large-scale production use.

4. Risks and Outlook

Single maintainer: Core code is still led by Jan Pfeifer alone; if his attention shifts, the project could slow down.
Model library gap: SOTA architectures like vision transformers, BERT, and diffusion models need community porting. The author plans to support HuggingFace ONNX import first, then native weight loading.
API stability: The official docs state "early stage; APIs may change slightly between versions" — watch for breaking changes when upgrading.

Conclusion

GoMLX is currently the only MLX-style frontend in the Go ecosystem that can run end-to-end CUDA training. The 2024–2025 updates moved it from "can run" to "barely usable," making it suitable for research, teaching, or edge scenarios where a single binary delivery matters. If you need ready-to-use HuggingFace models or large-scale distributed training, Python remains the better choice for now, with Go serving as glue for services. Otherwise, GoMLX is worth trying — and keep an eye on whether v0.5 delivers direct HuggingFace model import and multi-node parallelism.