MLX-Go

Go bindings for Apple MLX — a high-performance machine learning framework designed for Apple Silicon.

Note: This is a hobby project and public release of v0.1.0. Expect problems, API changes, and occasional breakage between releases while the library is still maturing.

Why MLX-Go?

• 🚀 Native Apple Silicon Performance — Leverage Metal GPU acceleration for ML workloads
• 🔧 Go Ecosystem Integration — Build inference servers with Go's concurrency primitives
• 📦 Simple, Familiar API — NumPy-like operations with idiomatic Go patterns
• 🧠 ML-Ready Operations — Comprehensive support for transformers, attention mechanisms, and more
• ⚡ Lazy Evaluation — Efficient computation graph optimization
• 🎯 Early-Stage but Ambitious — Strong foundations, but still evolving quickly

Quick Start

package main

import (
	mx "github.com/guptaaryan16/mlx_go/core"
)

func main() {
	// Create arrays
	a := mx.FromSlice([]float32{1, 2, 3, 4})
	b := mx.FromSlice([]float32{5, 6, 7, 8})
	
	// Operations are lazily evaluated
	c := mx.Add(a, b)
	d := mx.Multiply(c, c)
	
	// Trigger computation
	d.Eval()
	d.PrintArray()  // Output: [36, 64, 100, 144]
}

Installation

MLX-Go now uses a one-time native install plus ordinary Go module usage.

Supported targets

• darwin/arm64 with Apple Silicon and Apple frameworks
• linux/amd64 with system BLAS/LAPACK and the usual C++ runtime libraries

1. Install native MLX artifacts once per machine

go run github.com/guptaaryan16/mlx_go/cmd/mlxgo-install@v0.1.0

The installer downloads the native MLX release asset for your platform, installs the headers and static libraries into the standard local prefix, verifies the installation, and prints the exact next-step commands.

On Linux, installing into /usr/local may require sudo.

2. Use the Go module normally

go get github.com/guptaaryan16/mlx_go@v0.1.0

Then import it in your application:

import mx "github.com/guptaaryan16/mlx_go/core"

And build as usual:

go build ./...

3. Source-build fallback

If you need a custom prefix or an unsupported environment, build and install the native MLX libraries from source:

git clone --recursive https://github.com/guptaaryan16/mlx_go
cd mlx_go
./scripts/install.sh

If you are modifying csrc/mlx or csrc/mlx-c, build contributor-local native artifacts and use the explicit vendor mode:

./scripts/build_local.sh
go test -tags mlx_vendor ./...

For detailed installation instructions, see docs/installation.md.

Features

Core Operations

• Array Creation: Zeros, Ones, Eye, Arange, Linspace, FromSlice
• Element-wise Math: Add, Multiply, Exp, Log, Sqrt, Power, Sin, Cos, Tanh
• Reductions: Sum, Mean, Max, Min, Prod, Std, Var
• Broadcasting: Automatic shape broadcasting following NumPy semantics
• Reshaping: Reshape, Flatten, Squeeze, ExpandDims, Transpose
• Indexing: Advanced NumPy-style slicing and indexing

Neural Network Layers

• Activations: ReLU, GELU, SiLU, Sigmoid, Softmax, LogSoftmax
• Normalization: RMSNorm, LayerNorm
• Attention: Scaled Dot-Product Attention with causal masking support
• Position Encoding: Rotary Position Embeddings (RoPE)
• Layers: Linear, Embedding, Dropout
• Convolutions: Conv1d, Conv2d (via C API)

Linear Algebra

• Decompositions: QR, SVD, LU, Cholesky, Eigendecomposition
• Solvers: Linear systems, triangular solvers
• Operations: Matrix multiplication, norms, inverse, pseudo-inverse, trace

Random Operations

• Distributions: Uniform, Normal, Truncated Normal, Gumbel, Laplace
• Sampling: Categorical, Permutation, Randint
• Seeding: Reproducible random number generation

Documentation

Resource	Description
Installation Guide	Installer flow, source fallback, troubleshooting
Usage Guide	Comprehensive API overview with examples
Examples	Working code samples for common tasks
API Reference	Full API documentation on pkg.go.dev

Examples

Check out the examples/ directory for complete working examples:

• Linear Regression — Training a simple linear model
• Autograd Basics — Grad, ValueAndGrad, JVP, VJP, and Checkpoint walkthroughs
• Compiled Training — JIT-compiled linear regression training step
• BERT — Tiny BERT demo plus safetensors weight loading
• ResNet — Forward pass and optimizer-driven training loop
• Save/Load — Serializing and loading array data

More examples coming soon: LLaMA inference, custom layers, distributed serving.

Benchmarks

MLX-Go includes a comprehensive benchmarking tool to compare performance against Python MLX:

cd benchmark
go build -o benchmark_tool .
./benchmark_tool --quick

See benchmark/README.md for details.

Project Goals

MLX-Go aims to bring production-grade machine learning capabilities to the Go ecosystem:

1. Performance: Match or exceed Python MLX performance through efficient CGO usage
2. Ergonomics: Provide an idiomatic Go API that feels natural to Go developers
3. Completeness: Support the full range of MLX operations needed for inference and training
4. Reliability: Maintain strong type safety and comprehensive error handling

Read more about the project vision in GOAL.md.

Contributing

We welcome contributions! Please see CONTRIBUTING.md for guidelines on:

• Setting up your development environment
• Running tests and benchmarks
• Code style and best practices
• Submitting pull requests

Roadmap

• Core array operations and broadcasting
• Neural network primitives (layers, activations, attention)
• Linear algebra operations
• Comprehensive test coverage
• Training support (autograd, optimizers)
• Quantization (INT8, INT4 support)
• Model implementations (LLaMA, Whisper, CLIP)
• Distributed training support
• Performance optimizations (operation fusion, memory pooling)

License

This project is licensed under the MIT License - see the LICENSE file for details.

MLX-Go is built on top of Apple's MLX and MLX-C libraries.

Acknowledgments

• Apple's ML Explore team for creating MLX
• The Go community for excellent tooling and ecosystem
• I don't like to admit but Google's Antigravity helped me a lot for the project

---

Star this repo ⭐ if you find it useful! Contributions and feedback are always welcome.