mlx

README

MLX

Quickstart | Installation | Documentation | Examples

MLX is an array framework for machine learning on Apple silicon, brought to you by Apple machine learning research.

Some key features of MLX include:

• Familiar APIs: MLX has a Python API that closely follows NumPy. MLX also has fully featured C++, C, Swift, and Go APIs, which closely mirror the Python API. MLX has higher-level packages like mlx.nn and mlx.optimizers with APIs that closely follow PyTorch to simplify building more complex models.

• Composable function transformations: MLX supports composable function transformations for automatic differentiation, automatic vectorization, and computation graph optimization.

• Lazy computation: Computations in MLX are lazy. Arrays are only materialized when needed.

• Dynamic graph construction: Computation graphs in MLX are constructed dynamically. Changing the shapes of function arguments does not trigger slow compilations, and debugging is simple and intuitive.

• Multi-device: Operations can run on any of the supported devices (currently the CPU and the GPU).

• Unified memory: A notable difference from MLX and other frameworks is the unified memory model. Arrays in MLX live in shared memory. Operations on MLX arrays can be performed on any of the supported device types without transferring data.

MLX is designed by machine learning researchers for machine learning researchers. The framework is intended to be user-friendly, but still efficient to train and deploy models. The design of the framework itself is also conceptually simple. We intend to make it easy for researchers to extend and improve MLX with the goal of quickly exploring new ideas.

The design of MLX is inspired by frameworks like NumPy, PyTorch, Jax, and ArrayFire.

Examples

The MLX examples repo has a variety of examples, including:

• Transformer language model training.
• Large-scale text generation with LLaMA and finetuning with LoRA.
• Generating images with Stable Diffusion.
• Speech recognition with OpenAI's Whisper.

Quickstart

See the quick start guide in the documentation.

Installation

Python

MLX is available on PyPI. To install MLX on macOS, run:

pip install mlx

To install the CUDA backend on Linux, run:

pip install mlx[cuda]

To install a CPU-only Linux package, run:

pip install mlx[cpu]

Go

MLX has official Go bindings for high-performance computing. To install:

go get github.com/luxfi/mlx@latest

Or with specific version:

go get github.com/luxfi/mlx@v0.1.0

For GPU support, ensure CGO is enabled:

CGO_ENABLED=1 go build

See the Go documentation for detailed usage.

Checkout the documentation for more information on building the C++ and Python APIs from source.

Contributing

Check out the contribution guidelines for more information on contributing to MLX. See the docs for more information on building from source, and running tests.

We are grateful for all of our contributors. If you contribute to MLX and wish to be acknowledged, please add your name to the list in your pull request.

Citing MLX

The MLX software suite was initially developed with equal contribution by Awni Hannun, Jagrit Digani, Angelos Katharopoulos, and Ronan Collobert. If you find MLX useful in your research and wish to cite it, please use the following BibTex entry:

@software{mlx2023,
  author = {Awni Hannun and Jagrit Digani and Angelos Katharopoulos and Ronan Collobert},
  title = {{MLX}: Efficient and flexible machine learning on Apple silicon},
  url = {https://github.com/ml-explore},
  version = {0.0},
  year = {2023},
}

Documentation

Overview

Package mlx provides Go bindings for the MLX machine learning framework. MLX is an array framework for machine learning with multiple backend support:

- Metal: Apple Silicon GPUs (macOS) - CUDA: NVIDIA GPUs (Linux/Windows) - CPU: Fallback with SIMD optimizations (all platforms)

The package automatically detects and uses the best available backend. Build with CGO_ENABLED=1 and appropriate build tags for GPU support:

- cuda: Enable CUDA backend (requires CUDA toolkit) - metal: Enable Metal backend (automatic on macOS)

Index

• Variables
• func Eval(arrays ...*Array)
• func Info() string
• func SetBackend(backend Backend) error
• func Synchronize()
• type Array
  • func Add(a, b *Array) *Array
  • func Arange(start, stop, step float64) *Array
  • func ArrayFromSlice[T int64 | float64 | float32 | int32](data []T, shape []int, dtype Dtype) *Array
  • func FromSlice(data []float32, shape []int, dtype Dtype) *Array
  • func MatMul(a, b *Array) *Array
  • func Maximum(a, b *Array) *Array
  • func Mean(a *Array, axis ...int) *Array
  • func Multiply(a, b *Array) *Array
  • func Ones(shape []int, dtype Dtype) *Array
  • func Random(shape []int, dtype Dtype) *Array
  • func Sum(a *Array, axis ...int) *Array
  • func Zeros(shape []int, dtype Dtype) *Array
  • func (a *Array) Shape() []int
• type Backend
  • func GetBackend() Backend
  • func (b Backend) String() string
• type CPUDevice
  • func InitCPU() (*CPUDevice, error)
  • func (d *CPUDevice) Add(a, b []float32) ([]float32, error)
  • func (d *CPUDevice) MatMul(a, b []float32, m, n, k int) ([]float32, error)
• type Context
  • func (c *Context) Add(a, b *Array) *Array
  • func (c *Context) Arange(start, stop, step float64) *Array
  • func (c *Context) ArrayFromSlice(data any, shape []int, dtype Dtype) *Array
  • func (c *Context) Eval(arrays ...*Array)
  • func (c *Context) Free(a *Array)
  • func (c *Context) FreeStream(s *Stream)
  • func (c *Context) FromSlice(data []float32, shape []int, dtype Dtype) *Array
  • func (c *Context) GetBackend() Backend
  • func (c *Context) GetDevice() *Device
  • func (c *Context) MatMul(a, b *Array) *Array
  • func (c *Context) Maximum(a, b *Array) *Array
  • func (c *Context) Mean(a *Array, axis ...int) *Array
  • func (c *Context) Multiply(a, b *Array) *Array
  • func (c *Context) NewStream() *Stream
  • func (c *Context) Ones(shape []int, dtype Dtype) *Array
  • func (c *Context) Random(shape []int, dtype Dtype) *Array
  • func (c *Context) SetBackend(backend Backend) error
  • func (c *Context) Sum(a *Array, axis ...int) *Array
  • func (c *Context) Synchronize()
  • func (c *Context) Zeros(shape []int, dtype Dtype) *Array
• type Device
  • func GetDevice() *Device
• type DeviceInfo
• type Dtype
• type Stream
  • func NewStream() *Stream

Variables

var (
	// DefaultContext is the global MLX context
	DefaultContext *Context

	// ErrNoGPU is returned when GPU is requested but not available
	ErrNoGPU = errors.New("no GPU available")

	// ErrInvalidBackend is returned for invalid backend selection
	ErrInvalidBackend = errors.New("invalid backend")

	// Version is the MLX library version
	Version = "0.1.0"
)

Functions

func Eval

func Eval(arrays ...*Array)

Eval forces evaluation of lazy operations

func Info

func Info() string

Info returns information about the MLX installation

func SetBackend

func SetBackend(backend Backend) error

SetBackend sets the compute backend

func Synchronize

func Synchronize()

Synchronize waits for all operations to complete

Types

type Array

type Array struct {
	// contains filtered or unexported fields
}

Array represents a multi-dimensional array

func Add

func Add(a, b *Array) *Array

Add performs element-wise addition

func Arange

func Arange(start, stop, step float64) *Array

Arange creates an array with sequential values

func ArrayFromSlice

func ArrayFromSlice[T int64 | float64 | float32 | int32](data []T, shape []int, dtype Dtype) *Array

ArrayFromSlice creates an array from a typed Go slice with specified shape and dtype. Supports []int64, []float64, []float32, []int32 data types.

func FromSlice

func FromSlice(data []float32, shape []int, dtype Dtype) *Array

FromSlice creates an array from a Go slice

func MatMul

func MatMul(a, b *Array) *Array

MatMul performs matrix multiplication

func Maximum

func Maximum(a, b *Array) *Array

Maximum computes element-wise maximum of two arrays

func Mean

func Mean(a *Array, axis ...int) *Array

Mean computes the mean of array elements

func Multiply

func Multiply(a, b *Array) *Array

Multiply performs element-wise multiplication

func Ones

func Ones(shape []int, dtype Dtype) *Array

Ones creates an array filled with ones

func Random

func Random(shape []int, dtype Dtype) *Array

Random creates an array with random values

func Sum

func Sum(a *Array, axis ...int) *Array

Sum computes the sum of array elements

func Zeros

func Zeros(shape []int, dtype Dtype) *Array

Zeros creates a zero-filled array

func (*Array) Shape

func (a *Array) Shape() []int

Shape returns the shape of the array

type Backend

type Backend int

Backend represents the compute backend used by MLX

const (
	// CPU uses CPU-only computation
	CPU Backend = iota
	// Metal uses Apple Metal GPU acceleration
	Metal
	// CUDA uses NVIDIA CUDA GPU acceleration
	CUDA
	// ONNX uses ONNX Runtime (Windows fallback)
	ONNX
	// Auto automatically selects the best available backend
	Auto
)

func GetBackend

func GetBackend() Backend

GetBackend returns the current compute backend

func (Backend) String

func (b Backend) String() string

String returns the string representation of the backend

type CPUDevice

type CPUDevice struct {
	// contains filtered or unexported fields
}

CPUDevice represents CPU-based computation fallback

func InitCPU

func InitCPU() (*CPUDevice, error)

InitCPU initializes CPU-based computation

func (*CPUDevice) Add

func (d *CPUDevice) Add(a, b []float32) ([]float32, error)

Add performs element-wise addition on CPU

func (*CPUDevice) MatMul

func (d *CPUDevice) MatMul(a, b []float32, m, n, k int) ([]float32, error)

MatMul performs matrix multiplication on CPU

type Context

type Context struct {
	// contains filtered or unexported fields
}

Context manages MLX runtime and resources

func (*Context) Add

func (c *Context) Add(a, b *Array) *Array

Add performs element-wise addition

func (*Context) Arange

func (c *Context) Arange(start, stop, step float64) *Array

Arange creates an array with sequential values

func (*Context) ArrayFromSlice

func (c *Context) ArrayFromSlice(data any, shape []int, dtype Dtype) *Array

ArrayFromSlice creates an array from a typed slice (Context method)

func (*Context) Eval

func (c *Context) Eval(arrays ...*Array)

Eval forces evaluation of lazy operations

func (*Context) Free

func (c *Context) Free(a *Array)

Free releases resources for an array

func (*Context) FreeStream

func (c *Context) FreeStream(s *Stream)

FreeStream releases resources for a stream

func (*Context) FromSlice

func (c *Context) FromSlice(data []float32, shape []int, dtype Dtype) *Array

FromSlice creates an array from a Go slice

func (*Context) GetBackend

func (c *Context) GetBackend() Backend

GetBackend returns the current backend for this context

func (*Context) GetDevice

func (c *Context) GetDevice() *Device

GetDevice returns the current device for this context

func (*Context) MatMul

func (c *Context) MatMul(a, b *Array) *Array

MatMul performs matrix multiplication

func (*Context) Maximum

func (c *Context) Maximum(a, b *Array) *Array

Maximum computes element-wise maximum of two arrays

func (*Context) Mean

func (c *Context) Mean(a *Array, axis ...int) *Array

Mean computes the mean of array elements

func (*Context) Multiply

func (c *Context) Multiply(a, b *Array) *Array

Multiply performs element-wise multiplication

func (*Context) NewStream

func (c *Context) NewStream() *Stream

NewStream creates a new compute stream

func (*Context) Ones

func (c *Context) Ones(shape []int, dtype Dtype) *Array

Ones creates an array filled with ones

func (*Context) Random

func (c *Context) Random(shape []int, dtype Dtype) *Array

Random creates an array with random values

func (*Context) SetBackend

func (c *Context) SetBackend(backend Backend) error

SetBackend sets the compute backend for this context

func (*Context) Sum

func (c *Context) Sum(a *Array, axis ...int) *Array

Sum computes the sum of array elements

func (*Context) Synchronize

func (c *Context) Synchronize()

Synchronize waits for all operations to complete

func (*Context) Zeros

func (c *Context) Zeros(shape []int, dtype Dtype) *Array

Zeros creates a zero-filled array

type Device

type Device struct {
	Type   Backend
	ID     int
	Name   string
	Memory int64 // Memory in bytes
}

Device represents a compute device (CPU or GPU)

func GetDevice

func GetDevice() *Device

GetDevice returns the current compute device

type DeviceInfo

type DeviceInfo = Device

DeviceInfo is an alias for Device (for compatibility)

type Dtype

type Dtype int

Dtype represents the data type of array elements

const (
	Float32 Dtype = iota
	Float64
	Int32
	Int64
	Bool
)

type Stream

type Stream struct {
	// contains filtered or unexported fields
}

Stream represents a compute stream for async operations

func NewStream

func NewStream() *Stream

NewStream creates a new compute stream