mlx

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

# 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

# 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

# train.py
learning_rate = 0.001
batch_size = 32
num_epochs = 10
model = SimpleCNN(num_classes=10)

✅ GOOD: Config-Driven

# 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"

# 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.