gomlx/gomlx

GoMLX: An Accelerated Machine Learning Framework For Go

What it solves

GoMLX provides a comprehensive machine learning and mathematics framework for the Go programming language, serving as a Go-native alternative to popular frameworks like PyTorch, JAX, or TensorFlow. It enables developers to train, fine-tune, and execute machine learning models within the Go ecosystem, reducing the need to rely on Python for ML development and production.

How it works

GoMLX acts as an intermediary API that provides a library of ML layers and tools, which it then executes via interchangeable backends:

  • XLA Backend: Uses OpenXLA for high-performance JIT compilation to CPUs, GPUs (Nvidia, AMD ROCm, Intel, Macs), and TPUs. It supports distributed execution across multiple accelerators.
  • Go Backend: A pure Go implementation that is highly portable, supporting WASM for browser-based execution and optimized with SIMD for certain architectures.
  • ONNX Backend: Leverages ONNX Runtime to execute computations and allows models to be saved in the .onnx format.

The framework includes automatic differentiation (autodiff), a system for variable management (Store and Scope), and a wide array of differentiable operators and ML layers.

Who it’s for

It is designed for Go developers who want to build, experiment with, and productionize machine learning models without leaving the Go ecosystem, as well as researchers experimenting with non-conventional optimizer or regularizer ideas.

Highlights

  • HuggingFace Integration: Compatibility with go-huggingface for downloading models, tokenizers, and datasets, and reading parameters from GGUF or safetensors.
  • Diverse Layer Support: Includes FFN, convolutions, Multi-Head-Attention, LSTM, and specialized layers like KAN (Kolmogorov-Arnold Networks) and VNN (Vector Neural Networks).
  • Flexible Execution: Supports a variety of hardware targets via XLA, including multi-GPU/TPU distributed execution.
  • Developer Tooling: Includes a command-line tool for inspecting checkpoints and generating plots, and integration with JupyterLab via GoNB.
  • Advanced Training Features: Supports gradient checkpointing to trade memory for recomputation during large model training.

Related

  • Project
    yalue/onnxruntime_go

    A cross-platform Go wrapper for the onnxruntime library, allowing Go developers to load and execute ONNX-format neural networks with hardware acceleration support.

  • Project
    knights-analytics/hugot

    A Go library for running and fine-tuning ONNX transformer pipelines, enabling Hugging Face compatible inference and training natively in Golang.

  • Project
    microsoft/onnxruntime

    A cross-platform machine learning accelerator that optimizes inference and training for models from various frameworks like PyTorch and TensorFlow.

  • Project
    ml-explore/mlx

    An array framework by Apple machine learning research optimized for machine learning on Apple silicon, featuring a unified memory model and familiar NumPy/PyTorch-like APIs.

  • Project
    ekzhang/jax-js

    A high-performance machine learning framework for the browser that brings JAX-style numerical computing, automatic differentiation, and GPU acceleration via WebGPU and Wasm.