argmax returns a wrong result with default specialization on iPhone 17 Pro (uint8 and int32 outputs); CPU-only is correct

[glenn-jocher]

Summary

A graph that ends in argmax returns a wrong result on an iPhone 17 Pro when the .aimodel is loaded with SpecializationOptions.default. The same asset with the same input is correct with .cpuOnly. There is no error or warning, the result is deterministic, and it happens for both a uint8 and an int32 output, so the fault is in the argmax lane rather than in the cast.

The model is the YOLO26n semantic segmentation model (19 classes) with logits.argmax(1).to(torch.uint8) appended, so the graph emits a (1, 640, 640) class map instead of (1, 19, 640, 640) logits. The same model exported without the argmax is correct under .default: taking the argmax of its logits in Swift agrees with .cpuOnly on 99.4% of pixels.

This is separate from #115 (wrong outputs from a topk + gather postprocess), though both are silent wrong answers from the lane .default selects on this device.

Environment

• iPhone 17 Pro (iPhone18,1), iOS 27.0
• Xcode 27.0 (27A266a), iOS 27.0 SDK
• Export host: macOS 26.7, Apple silicon, Python 3.12, coreai-core==1.0.0b2, coreai-torch==0.4.2, torch==2.14.0, ultralytics==8.4.155

Export

pip install ultralytics==8.4.155 coreai-torch==0.4.2

# control: logits output (1, 19, 640, 640) float16, correct on device
yolo export model=yolo26n-sem.pt format=coreai quantize=16 imgsz=640

# class map: let the exporter append `argmax(1).to(uint8)` for Core AI as it already does for Core ML
python - <<'PY'
import pathlib, ultralytics.engine.exporter as E
p = pathlib.Path(E.__file__)
p.write_text(p.read_text().replace('fmt in {"qnn", "coreml", "ascend"}', 'fmt in {"qnn", "coreml", "ascend", "coreai"}'))
PY
cp yolo26n-sem.pt yolo26n-sem-classmap.pt
yolo export model=yolo26n-sem-classmap.pt format=coreai quantize=16 imgsz=640  # (1, 640, 640) uint8, wrong on device

yolo26n-sem.pt downloads automatically. The appended wrapper is ClassMapModel: y.argmax(1).to(torch.uint8). The rest of the export path is torch.export.export → run_decompositions(coreai_torch.get_decomp_table()) → TorchConverter().add_exported_program(...) → to_coreai() → optimize() → save_asset(), with the model and example input in half precision. The control asset is also published as yolo26n-sem.aimodel.zip.

Reproduction

Bundle the .aimodel directories and any photograph (this run used bus.jpg) in an iOS app and call compare(url, image:) for each asset. It runs the same pixels through .cpuOnly and .default and counts the pixels whose class differs.

// Minimal Core AI reproducer: same .aimodel, same image, CPU-only vs default specialization.
import CoreAI
import CoreGraphics
import Foundation
import ImageIO

/// Decodes an image and stretches it to size x size as RGB CHW in 0-1.
func pixels(_ url: URL, size: Int) -> [Float] {
  let source = CGImageSourceCreateWithURL(url as CFURL, nil)!
  let image = CGImageSourceCreateImageAtIndex(source, 0, nil)!
  var rgba = [UInt8](repeating: 0, count: size * size * 4)
  let context = CGContext(
    data: &rgba, width: size, height: size, bitsPerComponent: 8, bytesPerRow: size * 4,
    space: CGColorSpaceCreateDeviceRGB(), bitmapInfo: CGImageAlphaInfo.noneSkipLast.rawValue)!
  context.draw(image, in: CGRect(x: 0, y: 0, width: size, height: size))
  var chw = [Float](repeating: 0, count: size * size * 3)
  for i in 0..<(size * size) {
    for c in 0..<3 { chw[c * size * size + i] = Float(rgba[i * 4 + c]) / 255 }
  }
  return chw
}

/// Runs the model and returns the per-pixel class map (argmax over classes when the model emits logits).
@available(iOS 27.0, macOS 27.0, *)
func classMap(_ url: URL, _ options: SpecializationOptions, _ input: [Float]) async throws -> [UInt8] {
  let model = try await AIModel(contentsOf: url, options: options)
  let function = try model.loadFunction(named: "main")!
  let name = function.descriptor.inputNames[0]
  guard case .ndArray(let descriptor)? = function.descriptor.inputDescriptor(of: name) else { fatalError() }
  let array = NDArray(scalars: input.map { Float16($0) }, shape: descriptor.shape)  // FP16 assets
  var outputs = try await function.run(inputs: [name: array])
  let out = outputs.remove(function.descriptor.outputNames[0])!.ndArray!
  if out.scalarType == .uint8 {  // (1, H, W) class map computed in the graph
    let span = out.view(as: UInt8.self).contiguousElements!
    return (0..<span.count).map { span[$0] }
  }
  if out.scalarType == .int32 {  // same class map as int32
    let span = out.view(as: Int32.self).contiguousElements!
    return (0..<span.count).map { UInt8(clamping: span[$0]) }
  }
  let span = out.view(as: Float16.self).contiguousElements!  // (1, C, H, W) logits
  let classes = out.shape[1]
  let plane = out.shape[2] * out.shape[3]
  return (0..<plane).map { i in
    UInt8((0..<classes).max { span[$0 * plane + i] < span[$1 * plane + i] }!)
  }
}

func histogram(_ map: [UInt8]) -> String {
  var counts = [UInt8: Int]()
  for value in map { counts[value, default: 0] += 1 }
  return counts.sorted { $0.key < $1.key }.map { "\($0.key):\($0.value)" }.joined(separator: " ")
}

@available(iOS 27.0, macOS 27.0, *)
func compare(_ url: URL, image: URL) async {
  do {
    let input = pixels(image, size: 640)
    let cpu = try await classMap(url, .cpuOnly, input)
    let def = try await classMap(url, .default, input)
    let differing = zip(cpu, def).filter { $0 != $1 }.count
    print("REPRO \(url.lastPathComponent): \(differing) of \(cpu.count) pixels differ between cpuOnly and default")
    print("REPRO   cpuOnly \(histogram(cpu))")
    print("REPRO   default \(histogram(def))")
  } catch { print("REPRO \(url.lastPathComponent) failed: \(error)") }
}

Observed

REPRO yolo26n-sem-classmap.aimodel: 198442 of 409600 pixels differ between cpuOnly and default
REPRO   cpuOnly 0:118651 1:7758 2:79336 3:34 4:96 5:710 6:6 7:4351 8:19623 11:60880 13:7255 15:110900
REPRO   default 0:264154 1:3911 2:39656 3:11 4:58 5:319 6:2 7:2183 8:9790 11:30368 13:4270 15:54878
REPRO yolo26n-sem-int32.aimodel: 294172 of 409600 pixels differ between cpuOnly and default
REPRO   cpuOnly 0:118651 1:7758 2:79336 3:34 4:96 5:710 6:6 7:4351 8:19623 11:60880 13:7255 15:110900
REPRO   default 0:263837 1:52 5:15 7:5 11:193 13:18 15:50 255:145430
REPRO yolo26n-sem.aimodel: 2401 of 409600 pixels differ between cpuOnly and default
REPRO   cpuOnly 0:118651 1:7758 2:79336 3:34 4:96 5:710 6:6 7:4351 8:19623 11:60880 13:7255 15:110900
REPRO   default 0:118724 1:7815 2:79350 3:22 4:120 5:630 6:5 7:4341 8:19565 11:60734 13:8531 15:109763

Each line is a class:pixel count histogram of the 640x640 class map.

• uint8 class map, .default: 48% of pixels differ from .cpuOnly. Every non-zero class has almost exactly half its .cpuOnly count (79336 → 39656, 110900 → 54878, 60880 → 30368, 19623 → 9790) and the remainder is counted as class 0, which looks as if about half of the output elements are never written. The numbers were identical across two launches.
• int32 class map, .default (same export with the wrapper's dtype set to torch.int32): 72% of pixels differ, and 145,430 pixels hold a value above 255, which is not a valid class index (the reproducer clamps it to 255 for the histogram).
• Logits control, .default: 0.6% of pixels differ, which is FP16 noise at class boundaries.
• .cpuOnly gives the same histogram for all three assets.

Expected

The same class map as .cpuOnly, or a load-time error or fallback if the selected lane cannot run argmax correctly.