MLX Swift Framework

MLX Swift is Apple's high-performance machine learning framework designed specifically for Apple Silicon. It provides NumPy-like array operations with lazy evaluation, automatic differentiation, and unified CPU/GPU memory.

When to Use This Skill

• Array operations on Apple Silicon (MLXArray)
• Building neural networks (MLXNN)
• Training models with automatic differentiation
• Custom Metal kernels via MLXFast
• Performance optimization with JIT compilation

Architecture Overview

MLXOptimizers (Adam, AdamW, SGD, etc.)
       ↓
MLXNN (Layers, Modules, Losses)
       ↓
MLX (Arrays, Ops, Transforms, FFT, Linalg, Random)
       ↓
Cmlx (C/C++ bindings, Metal GPU)

Key File Reference

Purpose	File Path
Core array	Source/MLX/MLXArray.swift
Operations	Source/MLX/Ops.swift
Transforms	Source/MLX/Transforms.swift
Factory methods	Source/MLX/Factory.swift
Neural layers	Source/MLXNN/*.swift
Optimizers	Source/MLXOptimizers/Optimizers.swift
Fast ops	Source/MLX/MLXFast.swift
Custom kernels	Source/MLX/MLXFastKernel.swift
Wired memory coordinator	Source/MLX/WiredMemory.swift
GPU working-set helper	Source/MLX/GPU+Metal.swift

Quick Start

Basic Array Creation

import MLX

// Create arrays
let a = MLXArray([1, 2, 3, 4])
let b = MLXArray(0 ..< 12, [3, 4])  // Shape [3, 4]
let c = MLXArray.zeros([2, 3])
let d = MLXArray.ones([4, 4], dtype: .float32)

// Random arrays (use MLXRandom namespace or free functions)
let uniform = MLXRandom.uniform(0.0 ..< 1.0, [3, 3])
let normal = MLXRandom.normal([100])

Array Properties

let array = MLXArray(0 ..< 12, [3, 4])
array.shape    // [3, 4]
array.ndim     // 2
array.size     // 12
array.dtype    // .int64
array.count    // 3 (first dimension)

Basic Operations

let a = MLXArray([1.0, 2.0, 3.0])
let b = MLXArray([4.0, 5.0, 6.0])