GoMLX is an accelerated machine learning and math framework for Go, aiming to provide a PyTorch/Jax-like experience for Go developers. It targets researchers and developers who want to build, train, and fine-tune ML models within the Go ecosystem, offering a familiar API with a focus on transparency and extensibility.

How It Works

GoMLX leverages OpenXLA/PJRT as its primary backend, enabling just-in-time compilation for CPU and GPU acceleration, mirroring the performance of Jax and TensorFlow. An experimental, more portable Go-based backend is also available, though it is significantly slower and has limited operator coverage. The framework includes automatic differentiation, context management for ML models, and a comprehensive library of ML layers, optimizers, and loss functions.

Quick Start & Requirements

• Installation: For XLA backend (Linux recommended): curl -sSf https://raw.githubusercontent.com/gomlx/gopjrt/main/cmd/install_linux_amd64.sh | bash. For CUDA support: curl -sSf https://raw.githubusercontent.com/gomlx/gopjrt/main/cmd/install_cuda.sh | bash. SimpleGo backend requires import _ "github.com/gomlx/gomlx/backends/simplego".
• Prerequisites: Linux for XLA backend. CUDA Toolkit for GPU support. hdf5-tools may be needed. macOS support is experimental and lags behind.
• Docker: docker pull janpfeifer/gomlx_jupyterlab:latest and docker run -it --rm -p 8888:8888 -v "${PWD}":/home/jupyter/work janpfeifer/gomlx_jupyterlab:latest for an environment with JupyterLab and GoNB.
• Resources: Docker image is ~5GB.
• Docs: Tutorial, Godoc, Blog Post.

Highlighted Details

• OpenXLA/PJRT backend for high-performance CPU/GPU acceleration.
• Experimental pure Go backend (SimpleGo) for portability (Mac/Windows/Wasm).
• Support for converting ONNX models via onnx-gomlx and Hugging Face integration.
• Includes a GoNB kernel for Jupyter notebooks, facilitating interactive development.

Maintenance & Community

• Actively developed by Jan Pfeifer.
• Community support via GitHub Issues and Discussions.

Licensing & Compatibility

• Apache License Version 2.0. Permissive for commercial use and closed-source linking.

Limitations & Caveats

The framework is considered experimental, with occasional breaking changes. macOS support is outdated and considered broken for pre-built binaries, requiring manual compilation. Jacobian support is not yet implemented, only gradients.