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

---

📦 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

1. 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.
2. Precomputed RoPE Harmonic Tables:
   • Eliminates transcendental runtime trigonometry (sin/cos) by indexing static FP16 lookup buffers.
3. 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.