ax-coreml

CoreML — On-Device Machine Learning

Quick Patterns

Basic Conversion (PyTorch to CoreML)

import coremltools as ct
import torch

model.eval()
traced = torch.jit.trace(model, example_input)
mlmodel = ct.convert(
    traced,
    inputs=[ct.TensorType(shape=example_input.shape)],
    minimum_deployment_target=ct.target.iOS18
)
mlmodel.save("MyModel.mlpackage")

Load and Predict (Swift)

// Async load (preferred)
let config = MLModelConfiguration()
config.computeUnits = .all  // .cpuOnly, .cpuAndGPU, .cpuAndNeuralEngine
let model = try await MLModel.load(contentsOf: url, configuration: config)

// Async prediction (thread-safe)
let output = try await model.prediction(from: input)

Post-Training Compression

from coremltools.optimize.coreml import OpPalettizerConfig, OptimizationConfig, palettize_weights

config = OpPalettizerConfig(mode="kmeans", nbits=4, granularity="per_grouped_channel", group_size=16)
compressed = palettize_weights(model, OptimizationConfig(global_config=config))

Stateful Model (KV-Cache)

let state = model.makeState()
let output = try model.prediction(from: input, using: state) // state updated in-place

MLTensor (iOS 18+)

let scores = MLTensor(shape: [1, vocab_size], scalars: logits)
let topK = scores.topK(k: 10)
let probs = (topK.values / temperature).softmax()
let sampled = probs.multinomial(numSamples: 1)
let result = await sampled.shapedArray(of: Int32.self) // materialize

---

Decision Tree

Need on-device ML?
|
+-- Text generation (simple prompts, structured output)?
|   -> Foundation Models (ax-foundation-models), not CoreML
|
+-- Custom trained model / fine-tuned LLM?
|   -> CoreML
|   |
|   +-- PyTorch model to convert?
|   |   -> Pattern: Basic Conversion
|   +-- Model too large for device?
|   |   -> Pattern: Compression (palettization > quantization > pruning)
|   +-- Transformer with KV-cache?
|   |   -> Pattern: Stateful Models
|   +-- Multiple LoRA adapters?
|   |   -> Pattern: Multi-Function Models
|   +-- Pipeline stitching between models?
|   |   -> Pattern: MLTensor
|   +-- Concurrent predictions needed?
|       -> Pattern: Async Prediction
|
+-- Issue / not working?
    +-- Model won't load?          -> Diagnostics: Load Failures
    +-- Slow inference?            -> Diagnostics: Performance
    +-- High memory?               -> Diagnostics: Memory
    +-- Accuracy lost after compress? -> Diagnostics: Compression
    +-- Conversion fails?          -> Diagnostics: Conversion