Skip to content

log_softmax produces -inf on Apple Neural Engine in fp16 when one class dominates #2728

Description

@Ashutosh0x

Problem

The PyTorch log_softmax converter produces -inf values on Apple Neural Engine (ANE) in fp16 when the input has a dominant class with large logit values. This silently corrupts the output of every classification model using log_softmax, F.log_softmax, or F.cross_entropy on ANE.

Root Cause

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

python res = mb.softmax(x=x, axis=axis) res = mb.log(x=res)

This is a naive log(softmax(x)) decomposition. While softmax itself uses max-shift stabilization internally, the resulting probabilities for non-dominant classes underflow to 0 in fp16 (any probability below ~6e-5). Then log(0) -> -inf.

In fp32, these tiny probabilities are representable (e.g., 1e-38), so CPU and GPU compute units are unaffected. The bug is specific to fp16 on ANE.

Reproduction

`python
import torch
import coremltools as ct
import numpy as np

class LogSoftmaxModel(torch.nn.Module):
def forward(self, x):
return torch.nn.functional.log_softmax(x, dim=-1)

model = LogSoftmaxModel().eval()

One dominant class at index 3 with large logit

x = torch.tensor([[0.0, 0.0, 0.0, 50.0, 0.0, 0.0, 0.0, 0.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)

Expected: [[-50., -50., -50., 0., -50., -50., -50., -50.]]

print('CoreML:', coreml_out)

Broken: [[-inf, -inf, -inf, 0., -inf, -inf, -inf, -inf]]

`

Expected vs Actual

Input logit PyTorch (correct) CoreML fp16 (broken)
0.0 (non-dominant) -50.0 -inf (WRONG)
50.0 (dominant) 0.0 0.0 (correct)

Fix

Use the standard numerically stable log-softmax decomposition:

log_softmax(x) = x - max(x) - log(sum(exp(x - max(x))))

By subtracting max(x) first, all exp() arguments are <= 0, so values are in (0, 1]. The log of the sum is computed directly, avoiding the underflow-prone intermediate softmax probabilities.

This is the formula used by:

  • PyTorch's own fused log_softmax CUDA kernel
  • coremltools' TensorFlow frontend for _softmax_cross_entropy_with_logits
  • JAX's jax.nn.log_softmax

Impact

Every classification model using nn.LogSoftmax, F.log_softmax, or F.cross_entropy on ANE with fp16 precision. This includes BERT, ResNet, ViT, and most models that compute cross-entropy loss or log-probability outputs.

Environment

  • coremltools version: 9.0 (main branch, commit as of 2026-05-29)
  • Affected compute unit: Neural Engine (fp16)
  • Unaffected: CPU, GPU

Related Issues

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions