logcumsumexp overflows on Apple Neural Engine in fp16 for inputs > ~11.09 (output collapses to inf then NaN)

[Ashutosh0x]

Problem

The PyTorch torch.logcumsumexp converter produces inf / NaN on Apple Neural Engine (ANE) in fp16 for input values above ~11.09. This is because the converter computes exp(x) on raw input without any stabilization, and exp(11.09) ~ 65,504 which is the fp16 maximum.

Root Cause

The current converter at converters/mil/frontend/torch/ops.py line 2230 computes:

exp = mb.exp(x=x)           # raw exp, no max-shift
cumsumexp = mb.cumsum(x=exp, axis=dim)
res = mb.log(x=cumsumexp)

exp(x) is computed on raw input with zero stabilization. For any x > ~11.09, this overflows to inf in fp16, and the cumulative sum then propagates infinity.

Reproduction

import torch
import coremltools as ct

class LogCumSumExpModel(torch.nn.Module):
    def forward(self, x):
        return torch.logcumsumexp(x, dim=-1)

model = LogCumSumExpModel().eval()
x = torch.tensor([[1.0, 5.0, 10.0, 12.0, 15.0, 20.0, 50.0]])

traced = torch.jit.trace(model, x)
mlmodel = ct.convert(traced,
    inputs=[ct.TensorType(shape=x.shape)],
    compute_precision=ct.precision.FLOAT16)

pytorch_out = model(x).detach().numpy()
coreml_out = list(mlmodel.predict({'x_1': x.numpy()}).values())[0]

print('PyTorch:', pytorch_out)
print('CoreML:', coreml_out)  # inf/NaN for positions with x > 11.09

Fix

Use the standard max-shift stabilization:

logcumsumexp(x) = max(x) + log(cumsum(exp(x - max(x))))

By subtracting the global max(x) first, all exp() arguments are <= 0, so values are in (0, 1]. This is the same pattern used in the logsumexp stable decomposition (PR #2726).

Note: The global max is used rather than a running (cumulative) max because MIL does not provide a cummax op. The global max is always >= the running max at every position, so exp(x_i - global_max) <= 1 for all i, guaranteeing no overflow. The trade-off is slightly more underflow for early positions when a much larger value appears later, but this does not affect correctness.

Impact

Models using torch.logcumsumexp -- CTC decoders, autoregressive attention mechanisms, sequential probability models.

Environment

• coremltools version: 9.0 (main branch)
• Affected compute unit: Neural Engine (fp16)
• Unaffected: CPU, GPU

Related Issues

• Channel-reduce logsumexp on Apple Neural Engine has a hard fp16 overflow at x ≈ 7.63 (output drops to 0) #2690 -- reduce_log_sum_exp fp16 overflow (same pattern, fixed in PR Fix logsumexp fp16 overflow on ANE via stable max-shift decomposition #2726)
• Softplus on Apple Neural Engine has a hard fp16 discontinuity at x ≈ 10.4 (output drops to 0) #2687 -- Softplus fp16 overflow (same class, fixed in PR Fix softplus and mish fp16 overflow on ANE via stable decomposition #2725)
• log_softmax produces -inf on Apple Neural Engine in fp16 when one class dominates #2728 -- log_softmax fp16 underflow (same class, fix in PR Fix log_softmax fp16 underflow and logcumsumexp fp16 overflow on ANE via stable decomposition #2727)