MLX: Apple Silicon ML Framework

What is MLX? Apple open-source ML framework optimized for Apple Silicon (M1/M2/M3), with NumPy-like API and unified memory architecture.

Key Features

Basic Usage

import mlx.core as mx

# Create arrays
a = mx.array([1, 2, 3])
b = mx.array([4, 5, 6])

# Operations (lazy until evaluated)
c = a + b
d = mx.sum(c)

# Force evaluation
mx.eval(d)
print(d)  # 21

Neural Networks

import mlx.nn as nn

class MLP(nn.Module):
    def __init__(self, in_dim, hidden_dim, out_dim):
        super().__init__()
        self.linear1 = nn.Linear(in_dim, hidden_dim)
        self.linear2 = nn.Linear(hidden_dim, out_dim)

    def __call__(self, x):
        x = nn.relu(self.linear1(x))
        return self.linear2(x)

model = MLP(768, 512, 10)

MLX LLM

from mlx_lm import load, generate

model, tokenizer = load("mlx-community/Llama-3.2-3B-Instruct")

prompt = "Explain quantum computing in simple terms"
response = generate(model, tokenizer, prompt=prompt, max_tokens=200)
print(response)

Converting Models

# Convert HuggingFace to MLX
python -m mlx_lm.convert --hf-path meta-llama/Llama-3.2-3B-Instruct
    -q --q-bits 4  # Quantize to 4-bit

Performance on Apple Silicon

Training with MLX

import mlx.optimizers as optim

optimizer = optim.Adam(learning_rate=1e-3)

def loss_fn(model, x, y):
    return mx.mean((model(x) - y) ** 2)

loss_and_grad = nn.value_and_grad(model, loss_fn)

for batch in dataloader:
    loss, grads = loss_and_grad(model, batch.x, batch.y)
    optimizer.update(model, grads)
    mx.eval(model.parameters(), optimizer.state)

Comparison to PyTorch

Best Practices

• Use for local Mac development
• Convert model weights from HuggingFace
• Quantize for faster inference
• Use lazy evaluation pattern
• Great for experimentation