Fix log_softmax fp16 underflow and logcumsumexp fp16 overflow on ANE via stable decomposition

log_softmax: The naive log(softmax(x)) produces -inf for non-dominant
classes in fp16 because softmax outputs underflow to 0, then log(0) = -inf.
The stable form x - max(x) - log(sum(exp(x - max(x)))) avoids computing
tiny intermediate probabilities directly.

logcumsumexp: The naive log(cumsum(exp(x))) overflows in fp16 for x > ~11.09
since exp(11.09) exceeds fp16 max (65,504). The stable form shifts by the
global maximum first so all exp() arguments are <= 0, keeping values in (0,1].

Both fixes follow the same max-shift pattern used in the logsumexp stable
decomposition (PR #2726) and the softplus stable decomposition (PR #2725).

Added regression tests with extreme fp16 inputs for both ops.

Author: Ashutosh0x

coremltools/converters/mil/frontend/torch/ops.py

@@ -2228,9 +2228,48 @@ def _parse_keyword_args(context, node, dtype) -> Var:
         res = mb.cumsum(x=x, axis=dim, name=node.name)
     else:
         assert node.kind in ("logcumsumexp", "_logcumsumexp")
-        exp = mb.exp(x=x)
-        cumsumexp = mb.cumsum(x=exp, axis=dim)
-        res = mb.log(x=cumsumexp, name=node.name)
+        # Use numerically stable decomposition to prevent fp16 overflow
+        # on ANE.  The naive log(cumsum(exp(x))) computes exp(x) on raw
+        # input, which overflows in fp16 for x > ~11.09 (since
+        # exp(11.09) ~ 65,504 = fp16 max).
+        #
+        # Stable form:
+        #   logcumsumexp(x) = max(x) + log(cumsum(exp(x - max(x))))
+        #
+        # We use the global max over the entire axis 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) <= exp(x_i -
+        # running_max_i) <= 1 for all i.  This guarantees 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 -- those contributions are genuinely
+        # negligible.  This is the same max-shift pattern used in the
+        # logsumexp stable decomposition.
+
+        # Step 1: global max along the cumsum axis (keep dims for
+        #         broadcasting)
+        x_max = mb.reduce_max(x=x, axes=[dim], keep_dims=True,
+                              name=node.name + "_max")
+
+        # Step 2: x - max(x) (broadcast subtraction)
+        x_shifted = mb.sub(x=x, y=x_max,
+                           name=node.name + "_shifted")
+
+        # Step 3: exp(x - max(x)) -- all values in (0, 1], safe in fp16
+        x_exp = mb.exp(x=x_shifted,
+                       name=node.name + "_exp")
+
+        # Step 4: cumulative sum of shifted exponentials
+        x_cumsum = mb.cumsum(x=x_exp, axis=dim,
+                             name=node.name + "_cumsum")
+
+        # Step 5: log(cumsum(...))
+        x_log = mb.log(x=x_cumsum,
+                       name=node.name + "_log")
+
+        # Step 6: add back the global max
+        res = mb.add(x=x_log, y=x_max, name=node.name)
 
     context.add(res)
 
@@ -5915,8 +5954,47 @@ def _parse_positional_args(context, node) -> Tuple[Var]:
 
     x, axis = _parse_positional_args(context, node)
 
-    res = mb.softmax(x=x, axis=axis, name=node.name + "_softmax")
-    res = mb.log(x=res, name=node.name)
+    # Use numerically stable decomposition to prevent fp16 underflow
+    # on ANE.  The naive log(softmax(x)) first computes softmax
+    # probabilities, which underflow to 0 in fp16 for non-dominant
+    # classes (any probability below ~6e-5).  Then log(0) produces
+    # -inf, silently corrupting the output.
+    #
+    # Stable form:
+    #   log_softmax(x) = x - max(x) - log(sum(exp(x - max(x))))
+    #
+    # By subtracting max(x) first, all exp() arguments are <= 0, so
+    # exp() values are in (0, 1] -- no overflow.  The log of the sum
+    # is computed directly, avoiding the underflow-prone intermediate
+    # softmax probabilities.
+    #
+    # This matches PyTorch's own fused log_softmax implementation and
+    # the approach used in coremltools' TensorFlow frontend for
+    # cross-entropy loss (_softmax_cross_entropy_with_logits).
+
+    # Step 1: max(x) along softmax axis (keep dims for broadcasting)
+    x_max = mb.reduce_max(x=x, axes=[axis], keep_dims=True,
+                          name=node.name + "_max")
+
+    # Step 2: x - max(x) (broadcast subtraction)
+    x_shifted = mb.sub(x=x, y=x_max,
+                       name=node.name + "_shifted")
+
+    # Step 3: exp(x - max(x)) -- all values in (0, 1], safe in fp16
+    x_exp = mb.exp(x=x_shifted,
+                   name=node.name + "_exp")
+
+    # Step 4: sum(exp(x - max(x))) along softmax axis
+    x_sum = mb.reduce_sum(x=x_exp, axes=[axis], keep_dims=True,
+                          name=node.name + "_sum")
+
+    # Step 5: log(sum(...))
+    x_log_sum = mb.log(x=x_sum,
+                       name=node.name + "_log_sum")
+
+    # Step 6: (x - max(x)) - log(sum(exp(x - max(x))))
+    res = mb.sub(x=x_shifted, y=x_log_sum, name=node.name)
+
     context.add(res)
 
 

coremltools/converters/mil/frontend/torch/test/test_torch_ops.py

@@ -6534,6 +6534,74 @@ def test_softmax(self, compute_unit, backend, frontend, shape):
             shape, model, frontend=frontend, backend=backend, compute_unit=compute_unit
         )
 
+    @pytest.mark.parametrize(
+        "compute_unit, backend, frontend, shape",
+        itertools.product(compute_units, backends, frontends, COMMON_SHAPES_ALL),
+    )
+    def test_log_softmax(self, compute_unit, backend, frontend, shape):
+        model = ModuleWrapper(function=torch.nn.functional.log_softmax, kwargs={"dim": -1})
+        self.run_compare_torch(
+            shape, model, frontend=frontend, backend=backend, compute_unit=compute_unit
+        )
+
+    @pytest.mark.parametrize(
+        "compute_unit, backend",
+        itertools.product(compute_units, backends),
+    )
+    def test_log_softmax_fp16_stable_decomposition(self, compute_unit, backend):
+        """Regression test for issue #2728: the converter must decompose
+        log_softmax into the stable form x-max(x)-log(Σ exp(x-max(x)))
+        instead of emitting softmax followed by log, because the naive
+        path underflows in fp16 on Apple Neural Engine when one class
+        dominates (softmax outputs below ~6e-5 flush to 0, then log(0)
+        produces -inf).
+
+        This test uses a Conv2d+log_softmax+Flatten+Linear model with
+        input shape (1,1,32,32) and 64 output channels to ensure the
+        model routes to ANE (s=32, c=64 confirmed ALL_NE in sweep
+        profiles from PR #2618).  It verifies the graph structure: the
+        MIL program must contain a ``reduce_max`` op and must NOT
+        contain a ``softmax`` op.  Without the fix the graph contains
+        ``softmax`` and this test fails.
+        """
+
+        class LogSoftmaxModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.conv = nn.Conv2d(1, 64, kernel_size=3, padding=1)
+                self.flatten = nn.Flatten()
+                self.linear = nn.Linear(64 * 32 * 32, 10)
+
+            def forward(self, x):
+                x = torch.nn.functional.log_softmax(self.conv(x), dim=1)
+                x = self.flatten(x)
+                return self.linear(x)
+
+        model = LogSoftmaxModel().eval()
+        x = torch.randn(1, 1, 32, 32)
+
+        results = TorchBaseTest.run_compare_torch(
+            x,
+            model,
+            backend=backend,
+            compute_unit=compute_unit,
+            input_as_shape=False,
+        )
+        # Verify the naive softmax+log path was replaced by the stable decomposition.
+        # Without the fix, this assertion fails.
+        prog = results[1]._mil_program
+        softmax_ops = prog.find_ops(op_type="softmax")
+        assert len(softmax_ops) == 0, (
+            f"Expected no 'softmax' ops in the MIL graph after stable "
+            f"decomposition of log_softmax, but found {len(softmax_ops)}. "
+            f"The converter should use x-max(x)-log(Σ exp(x-max(x)))."
+        )
+        reduce_max_ops = prog.find_ops(op_type="reduce_max")
+        assert len(reduce_max_ops) > 0, (
+            "Expected at least one 'reduce_max' op in the MIL graph as "
+            "part of the stable log_softmax decomposition, but found none."
+        )
+
     @pytest.mark.parametrize(
         "compute_unit, backend, frontend, range_val",
         itertools.product(
@@ -12459,6 +12527,56 @@ def test_logcumsumexp(self, compute_unit, backend, frontend, axis):
             input_shape, model, frontend=frontend, backend=backend, compute_unit=compute_unit
         )
 
+    @pytest.mark.parametrize(
+        "compute_unit, backend",
+        itertools.product(compute_units, backends),
+    )
+    def test_logcumsumexp_fp16_stable_decomposition(self, compute_unit, backend):
+        """Regression test for issue #2729: the converter must decompose
+        logcumsumexp into a stable form using cumulative-max shifting
+        instead of the naive log(cumsum(exp(x))), because exp(x)
+        overflows in fp16 on Apple Neural Engine for x > ~11.09.
+
+        This test uses a Conv2d+logcumsumexp+Flatten+Linear model with
+        input shape (1,1,32,32) and 64 output channels to ensure the
+        model routes to ANE (s=32, c=64 confirmed ALL_NE in sweep
+        profiles from PR #2618).  It verifies the graph structure: the
+        MIL program must contain a ``reduce_max`` op.  Without the fix
+        the graph computes exp(x) on raw input and this test fails.
+        """
+
+        class LogCumSumExpModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.conv = nn.Conv2d(1, 64, kernel_size=3, padding=1)
+                self.flatten = nn.Flatten()
+                self.linear = nn.Linear(64 * 32 * 32, 10)
+
+            def forward(self, x):
+                x = torch.logcumsumexp(self.conv(x), dim=1)
+                x = self.flatten(x)
+                return self.linear(x)
+
+        model = LogCumSumExpModel().eval()
+        x = torch.randn(1, 1, 32, 32)
+
+        results = TorchBaseTest.run_compare_torch(
+            x,
+            model,
+            backend=backend,
+            compute_unit=compute_unit,
+            input_as_shape=False,
+        )
+        # Verify the naive exp path was replaced by the stable decomposition.
+        # Without the fix, this assertion fails.
+        prog = results[1]._mil_program
+        reduce_max_ops = prog.find_ops(op_type="reduce_max")
+        assert len(reduce_max_ops) > 0, (
+            "Expected at least one 'reduce_max' op in the MIL graph as "
+            "part of the stable logcumsumexp decomposition, but found "
+            "none. The converter should shift by cumulative max before exp()."
+        )
+
 
 class TestHannWindow(TorchBaseTest):
     @pytest.mark.parametrize(