mlx
This guide provides opinionated, actionable best practices for developing high-performance machine learning applications with MLX on Apple Silicon, focusing on Python and Swift.
Le skill
# mlx Best Practices
MLX is Apple's high-performance array framework for numerical computing and AI on Apple Silicon. Adhering to these guidelines ensures your MLX code is efficient, maintainable, and integrates seamlessly with Apple's ML ecosystem.
## 1. Code Organization and Structure
Organize your MLX projects into logical, reusable modules. Mirror the structure of official Apple examples (e.g., `corenet` MLX examples) for clarity and maintainability.
### 1.1 Modularize Components
Separate concerns into dedicated files: data loading, model definition, training loops, and utility functions.
❌ **BAD: Monolithic Script**
```python
# train.py
import mlx.core as mx
import mlx.nn as nn
# ... data loading, model definition, training loop all in one file ...
```
✅ **GOOD: Modular Structure**
```
my_project/
├── data/
│ └── mnist_loader.py
├── models/
│ └── simple_cnn.py
├── config/
│ └── train_config.yaml
├── utils/
│ └── metrics.py
└── train.py
```
```python
# models/simple_cnn.py
import mlx.core as mx
import mlx.nn as nn
class SimpleCNN(nn.Module):
def __init__(self, num_classes: int):
super().__init__()
self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
self.relu1 = nn.ReLU()
self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
self.linear = nn.Linear(32 * 14 * 14, num_classes)
def __call__(self, x: mx.array) -> mx.array:
x = self.conv1(x)
x = self.relu1(x)
x = self.pool1(x)
x = x.reshape(x.shape[0], -1) # Flatten
return self.linear(x)
# train.py
from data.mnist_loader import load_mnist_data
from models.simple_cnn import SimpleCNN
# ... training logic ...
```
### 1.2 Externalize Configuration
Store training parameters, model hyperparameters, and dataset paths in version-controlled YAML or JSON files.
❌ **BAD: Hardcoded Parameters**
```python
# train.py
learning_rate = 0.001
batch_size = 32
num_epochs = 10
model = SimpleCNN(num_classes=10)
```
✅ **GOOD: Config-Driven**
```yaml
# config/train_config.yaml
training:
learning_rate: 0.001
batch_size: 32
num_epochs: 10
model:
name: SimpleCNN
num_classes: 10
data:
path: "./data/mnist"
```
```python
# train.py
import yaml
from models.simple_cnn import SimpleCNN
with open("config/train_config.yaml", "r") as f:
config = yaml.safe_load(f)
model = SimpleCNN(num_classes=config["model"]["num_classes"])
# ... use config["training"]["learning_rate"] etc.
```
## 2. Common Patterns and Anti-patterns
Leverage MLX's core features like lazy evaluation and unified memory for optimal performance.
### 2.1 Embrace Lazy Evaluation
MLX computations are lazy; they build a computation graph but execute only when results are explicitly needed. Use `mlx.eval()` or `mlx.compile()` to force execution.
❌ **BAD: Relying on Implicit Evaluation for Performance Critical Sections**
```python
# This loop might build a huge graph or trigger many small evaluations
# if intermediate results are printed or accessed.
for i in range(num_steps):
loss, grads = value_and_grad(model, loss_fn)(model, X_batch, y_batch)
# ... optimizer update ...
# No explicit eval, performance can be unpredictable.
```
✅ **GOOD: Explicit Evaluation for Control**
Use `mlx.eval()` on the final outputs of a computation step to trigger execution efficiently. For static graphs, `mlx.compile()` offers further optimization.
```python
import mlx.core as mx
import mlx.nn as nn
from mlx.optimizers import Adam
from mlx.utils import value_and_grad
# Define model and loss_fn (as in 1.1)
def train_step(model, X, y, optimizer):
loss, grads = value_and_grad(model, loss_fn)(model, X, y)
optimizer.update(model, grads)
return loss
# Compile the training step for static graphs
compiled_train_step = mx.compile(train_step)
# In your training loop:
optimizer = Adam(learning_rate=0.001)
for i in range(num_steps):
X_batch, y_batch = get_batch()
# Use the compiled step
loss = compiled_train_step(model, X_batch, y_batch, optimizer)
# Explicitly evaluate the loss to trigger computation and get its value
mx.eval(model, optimizer, loss) # Evaluate model state, optimizer state, and loss
print(f"Step {i}, Loss: {loss.item()}")
```
### 2.2 Leverage Unified Memory
MLX arrays live in shared memory, eliminating explicit data transfers between CPU and GPU. Specify the device for operations when necessary, but default behavior is often sufficient.
❌ **BAD: Unnecessary Device Management (like other frameworks)**
```python
# No need to explicitly move data to 'gpu' or 'cpu'
data = mx.array([1, 2, 3]).to("gpu") # This is not how MLX works
```
✅ **GOOD: Device-Agnostic Operations (or explicit per-op)**
```python
import mlx.core as mx
a = mx.array([1, 2, 3])
b = mx.array([4, 5, 6])
# Operations run on the default device (usually GPU if available)
c = a + b
mx.eval(c)
# You can explicitly specify device per operation if needed for specific tasks
d = mx.add(a, b, stream=mx.cpu) # Force CPU execution for this specific op
mx.eval(d)
```
### 2.3 Use Type Hints and Pure Functions
Follow Python best practices: annotate functions and variables with type hints for clarity. Design functions to be pure where possible for easier testing and composability.
## 3. Performance Considerations
Optimize for Apple Silicon by understanding MLX's execution model.
### 3.1 Batching and Data Loading
Efficient data loading and proper batching are critical for GPU utilization.
❌ **BAD: Iterating over single samples or inefficient data loading**
```python
# Inefficient for GPU
for sample_x, sample_y in dataset:
# ... process single sample ...
```
✅ **GOOD: Batch Data and Preload**
```python
# Use a DataLoader or similar utility to yield batches of MLX arrays
def data_loader(data, batch_size):
for i in range(0, len(data), batch_size):
yield mx.array(data[i:i+batch_size])
# In training loop:
for X_batch, y_batch in data_loader(train_data, batch_size):
# ... process batch ...
```
### 3.2 Quantization for Deployment
For deploying models on-device, consider quantizing weights (e.g., 4-bit) to reduce memory footprint and improve inference speed.
```python
import mlx.core as mx
import mlx.nn as nn
# Assuming 'model' is a trained MLX model
# Quantize the model to 4-bit
quantized_model = model.quantize(group_size=32, bits=4)
# Save the quantized model
mx.save("quantized_model.safetensors", quantized_model.parameters())
```
## 4. Common Pitfalls and Gotchas
Avoid common mistakes that can lead to unexpected behavior or poor performance.
### 4.1 Forgetting `mlx.eval()`
The most common pitfall. If you don't explicitly evaluate, your graph might not execute, or only partially.
❌ **BAD: No `eval()` on results**
```python
# This will print the MLX array object, not its computed value
result = mx.add(mx.array([1]), mx.array([2]))
print(result) # Output: array(value=..., shape=(1,), dtype=int32)
```
✅ **GOOD: Always `eval()` when you need the value**
```python
result = mx.add(mx.array([1]), mx.array([2]))
mx.eval(result)
print(result) # Output: array(3, shape=(1,), dtype=int32)
print(result.item()) # Output: 3
```
### 4.2 Over-evaluating or Under-evaluating
Finding the right balance for `mlx.eval()` is key. Too many small `eval` calls can break graph optimization. Too few can lead to memory issues or delayed execution.
**Guideline:** Evaluate at logical synchronization points, typically at the end of a training step or inference call. Use `mx.eval(model, optimizer, loss)` to evaluate all relevant states.
## 5. Testing Approaches
Ensure the reliability and correctness of your MLX models.
### 5.1 Unit Tests for Model Components
Test individual layers, custom modules, and utility functions.
```python
# tests/test_simple_cnn.py
import mlx.core as mx
import mlx.nn as nn
from models.simple_cnn import SimpleCNN
def test_simple_cnn_output_shape():
model = SimpleCNN(num_classes=10)
# Initialize model parameters by running a dummy input
_ = model(mx.zeros((1, 1, 28, 28)))
mx.eval(model) # Ensure parameters are initialized
input_tensor = mx.random.normal((2, 1, 28, 28))
output = model(input_tensor)
assert output.shape == (2, 10)
print("SimpleCNN output shape test passed.")
if __name__ == "__main__":
test_simple_cnn_output_shape()
```
### 5.2 Integration Tests for Inference
Verify that your trained models perform as expected on new data, especially when integrating with Core ML or other on-device frameworks.
### 5.3 Data Validation and Reproducibility
Implement sanity checks for all input data. Version control your data, models, and training configurations to ensure reproducible results.
```python
# utils/data_validation.py
import mlx.core as mx
def validate_image_data(images: mx.array, expected_shape: tuple, expected_dtype: mx.Dtype):
assert images.ndim == len(expected_shape), f"Expected {len(expected_shape)} dims, got {images.ndim}"
assert images.shape == expected_shape, f"Expected shape {expected_shape}, got {images.shape}"
assert images.dtype == expected_dtype, f"Expected dtype {expected_dtype}, got {images.dtype}"
assert images.min() >= 0 and images.max() <= 1, "Image pixel values out of range [0, 1]"
print("Image data validation passed.")
# In your data loader or training script:
# ... load raw data ...
# raw_images = mx.array(...)
# validate_image_data(raw_images, (100, 1, 28, 28), mx.float32)
```
## 6. Swift-Specific Guidelines
When working with MLX Swift, adhere to Apple's official Swift style guides and API design guidelines.
### 6.1 Follow Swift API Style
Use idiomatic Swift patterns, consistent naming, and proper error handling. Refer to the `mlx-swift` examples (e.g., MNISTTrainer) as templates.
### 6.2 Manage Object Lifetime
Be mindful of object lifetimes, especially when bridging C++ MLX objects to Swift, as aggressive inlining and optimizations can affect behavior. The `mlx-swift` package includes recommendations for this.
Comment l’installer
1. Open the tool and start a new conversation, project or template. 2. Paste the text below wherever the tool keeps standing instructions — usually called system prompt, instructions, or persona. 3. Save it so it applies to every run rather than a single message.