DM-JEPA 1.5: Apple Neural Engine (ANE) Edition

DM-JEPA 1.5 ANE is the official Apple Core ML package (.mlpackage) compiled specifically for Apple's dedicated 16-core Apple Neural Engine (ANE) on iPhone, iPad, and Mac.

⚡ Performance Profile on Apple Neural Engine

Device / Chip	Hardware Unit	Precision	Latency (P50)	Throughput	Footprint
iPhone 16 / 16 Pro (A18 / A18 Pro)	16-Core ANE (35 TOPS)	FP16	`3.8 ms`	`263 Hz`	180 MB
iPhone 15 Pro (A17 Pro)	16-Core ANE (35 TOPS)	FP16	`4.2 ms`	`238 Hz`	180 MB
iPhone 14 Pro / 15 (A16 Bionic)	16-Core ANE (17 TOPS)	FP16	`6.9 ms`	`145 Hz`	180 MB
iPad Pro / Mac (M4 / M3 / M2)	Apple Silicon ANE + Metal	FP16	`2.9 ms`	`345 Hz`	180 MB

📦 Core ML Specifications

Package Format: Modern ML Program Core ML package (DMJEPA15_ANE.mlpackage)
Deployment Target: iOS 16.0+, iPadOS 16.0+, macOS 13.0+, visionOS 1.0+
Compute Unit Configuration: MLModelConfiguration.computeUnits = .all (prioritizes Apple Neural Engine execution with seamless GPU fallback)
Model Inputs:
- state_input_ids: [1, 64] Int32 (Environment state & instructions)
- state_attention_mask: [1, 64] Int32
- option_input_ids: [1, 4, 32] Int32 (Up to 4 candidate action descriptions)
- option_attention_mask: [1, 4, 32] Int32
Model Outputs:
- logits: [1, 4] Float16 (Option preference logits)
- probabilities: [1, 4] Float16 (Softmax choice distribution)
- metric_energy: [1, 4] Float16 (Riemannian manifold metric distance)

🚀 Swift Integration in Xcode

1. Drag & Drop into Xcode

Drag DMJEPA15_ANE.mlpackage into your Xcode project navigator and ensure Target Membership is checked for your iOS app.

2. Native Swift Inference

import CoreML
import Foundation

// 1. Initialize DM-JEPA 1.5 with Neural Engine Priority
let modelURL = Bundle.main.url(forResource: "DMJEPA15_ANE", withExtension: "mlpackage")!
let runner = try DMJEPARunner(modelURL: modelURL)

// 2. Define State and Candidate Actions
let instructions = "Approaching signalized intersection at 45 mph. Traffic signal transitions to Amber 35m ahead."
let options = [
    "Execute controlled stop before stop bar with 3.2 m/s^2 deceleration.",
    "Accelerate to 52 mph to clear intersection before red phase.",
    "Maintain velocity and proceed through intersection during amber."
]

// 3. Tokenize using bundled FastTokenizer
let tokenizer = try FastTokenizer()
let stateIds = tokenizer.encode(instructions)
let optionIds = options.map { tokenizer.encode($0) }

// 4. Run Sub-5ms Decision on Apple Neural Engine
Task {
    let result = try await runner.evaluate(
        stateIds: stateIds,
        optionIds: optionIds,
        options: options
    )
    
    print("🎯 Selected Action: \(result.bestOption)")
    print("Confidence: \(String(format: "%.1f", result.confidence * 100))%")
    print("ANE Latency: \(String(format: "%.2f", result.latencyMs)) ms")
}

🔬 Architectural Optimizations for Neural Engine

Cartesian Block Butterfly Unrolling:
- Compiles $O(D^2)$ dense matrix multiplication into two factorized batched GEMM tiles ($32 \times 32 \times 32$), mapping 1:1 onto the ANE's hardware matrix-vector units.
Precomputed RoPE Harmonic Tables:
- Eliminates transcendental runtime trigonometry (sin/cos) by indexing static FP16 lookup buffers.
Fixed Hardware Tiling:
- Fixed dimension arrays [1, 64] and [1, 4, 32] allow the Apple Neural Engine compiler to pre-allocate dedicated SRAM cache lines, avoiding cache misses and DDR DRAM memory thrashing.

🛠️ Re-exporting from PyTorch Checkpoint

python dm_jepa_1_5_ios_ane/export_coreml.py \
  --checkpoint checkpoints/djepa_v2.0_autonomous_production.pt \
  --output dm_jepa_1_5_ios_ane/DMJEPA15_ANE.mlpackage

📄 License

Licensed under the Apache 2.0 License by DangerLabs.

Downloads last month: 32