fix(torch-frontend): numerically stable log_softmax and logcumsumexp

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

@@ -2146,7 +2146,14 @@ def Sigmoid(context, node):
 @register_tf_op
 def Softplus(context, node):
     x = context[node.inputs[0]]
-    x = mb.softplus(x=x, name=node.name)
+    # Numerically stable softplus: max(x, 0) + log(1 + exp(-|x|)).
+    # Since -|x| <= 0, exp(-|x|) is always in (0, 1], so no overflow occurs
+    # in any precision.  This avoids fp16 overflow on the Apple Neural Engine
+    # at x > ~10.4 (see issues #2687, #2747).
+    abs_x = mb.abs(x=x)
+    exp_val = mb.exp(x=mb.mul(x=-1.0, y=abs_x))
+    log_val = mb.log(x=mb.add(x=1.0, y=exp_val))
+    x = mb.add(x=mb.maximum(x=x, y=0.0), y=log_val, name=node.name)
     context.add(node.name, x)
 
 

coremltools/converters/mil/frontend/tensorflow/test/test_ops.py

@@ -432,6 +432,45 @@ def build_model(x):
             backend=backend,
         )
 
+    @pytest.mark.parametrize(
+        "compute_unit, backend",
+        itertools.product(
+            [ct.ComputeUnit.CPU_ONLY, ct.ComputeUnit.ALL],
+            backends,
+        ),
+    )
+    def test_softplus_large_values(self, compute_unit, backend):
+        """Verify softplus produces correct output for large inputs (#2747).
+
+        The native mb.softplus op overflows in fp16 on some ANE hardware for
+        x > ~10.4 (exp(x) exceeds 65504). The stable decomposition
+        max(x, 0) + log(1 + exp(-|x|)) avoids this since exp(-|x|) is in (0, 1].
+
+        Note: On CPU, the native softplus produces correct output regardless of
+        this fix because CPU uses fp32. The overflow only manifests on the ANE
+        with fp16 compute. This test verifies the output is correct for large
+        values that would overflow naive fp16 softplus.
+        """
+        input_shape = (1, 6)
+
+        @make_tf_graph([input_shape])
+        def build_model(x):
+            return tf.math.softplus(x)
+
+        model, inputs, outputs = build_model
+        # Values chosen to overflow naive fp16 softplus: exp(15) ~ 3.3M, exp(20) ~ 485M.
+        # With the stable decomposition, exp(-|x|) is always in (0, 1], so no overflow.
+        input_values = [np.array([[15.0, 20.0, -5.0, 0.0, 1.0, 10.5]], dtype=np.float32)]
+        input_dict = dict(zip(inputs, input_values))
+
+        self.run_compare_tf(
+            model,
+            input_dict,
+            outputs,
+            compute_unit=compute_unit,
+            backend=backend,
+        )
+
     @pytest.mark.parametrize(
         "compute_unit, backend, rank_and_axes",
         itertools.product(

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

@@ -2235,9 +2235,23 @@ 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)
+        # Numerically stable logcumsumexp: shift by the global max along
+        # the cumulative dimension before exp, then shift back.
+        # The naive decomposition exp -> cumsum -> log overflows in fp16
+        # when exp(x) exceeds 65504 (x > ~10.4), producing inf/NaN on the
+        # Apple Neural Engine. By subtracting the max along the dim first,
+        # exp arguments are <= 0, preventing overflow.
+        # Note: this uses the global max rather than a cumulative max (which
+        # is not available as a MIL op), so early elements far below the max
+        # may still underflow to -inf in fp16. This is strictly better than
+        # the naive version which overflows to +inf/NaN for all elements.
+        # See issue #2729.
+        max_val = mb.reduce_max(x=x, axes=[dim.val], keep_dims=True, name=node.name + "_max")
+        shifted = mb.sub(x=x, y=max_val, name=node.name + "_shifted")
+        exp = mb.exp(x=shifted)
         cumsumexp = mb.cumsum(x=exp, axis=dim)
-        res = mb.log(x=cumsumexp, name=node.name)
+        log_cumsum = mb.log(x=cumsumexp)
+        res = mb.add(x=log_cumsum, y=max_val, name=node.name)
 
     context.add(res)
 
@@ -5922,8 +5936,16 @@ 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)
+    # Numerically stable log_softmax: x - reduce_log_sum_exp(x, axis).
+    # The naive decomposition softmax -> log overflows in fp16 when
+    # exp(x) exceeds 65504 (x > ~10.4), producing -inf for non-dominant
+    # classes on the Apple Neural Engine. Using reduce_log_sum_exp avoids
+    # this because it applies the max-shift trick internally:
+    #   log(sum(exp(x))) = max(x) + log(sum(exp(x - max(x))))
+    # This keeps exp arguments <= 0, preventing overflow.
+    # See issues #2728, #2690.
+    log_sum_exp = mb.reduce_log_sum_exp(x=x, axes=axis, keep_dims=True, name=node.name + "_logsumexp")
+    res = mb.sub(x=x, y=log_sum_exp, name=node.name)
     context.add(res)
 
 

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

@@ -6539,6 +6539,50 @@ 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 = nn.LogSoftmax().eval()
+        self.run_compare_torch(
+            shape, model, frontend=frontend, backend=backend, compute_unit=compute_unit
+        )
+
+    @pytest.mark.parametrize(
+        "compute_unit, backend, frontend",
+        itertools.product(compute_units, backends, frontends),
+    )
+    def test_log_softmax_large_values(self, compute_unit, backend, frontend):
+        """Verify log_softmax produces correct output for large inputs (#2728).
+
+        The naive decomposition softmax -> log overflows in fp16 on the
+        Apple Neural Engine when one class dominates: softmax probabilities
+        for non-dominant classes underflow to 0, producing -inf.
+        The stable decomposition x - reduce_log_sum_exp(x) avoids this.
+        """
+        input_shape = (2, 10, 4, 5)
+
+        class LogSoftmaxModel(nn.Module):
+            def __init__(self, dim):
+                super().__init__()
+                self.dim = dim
+
+            def forward(self, x):
+                return torch.log_softmax(x, dim=self.dim)
+
+        model = LogSoftmaxModel(dim=-1).eval()
+        # Values chosen to create dominant classes: scale by 15 so
+        # exp(15) ~ 3.3M, causing fp16 overflow in naive softmax.
+        input_values = torch.randn(*input_shape) * 15.0
+        self.run_compare_torch(
+            input_values,
+            model,
+            backend=backend,
+            compute_unit=compute_unit,
+            frontend=frontend,
+        )
+
     @pytest.mark.parametrize(
         "compute_unit, backend, frontend, range_val",
         itertools.product(
@@ -12568,6 +12612,43 @@ 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, frontend, axis",
+        itertools.product(
+            compute_units,
+            backends,
+            frontends,
+            [-1, 0, 1],
+        ),
+    )
+    def test_logcumsumexp_large_values(self, compute_unit, backend, frontend, axis):
+        """Verify logcumsumexp produces correct output for large inputs (#2729).
+
+        The naive decomposition exp -> cumsum -> log overflows in fp16 on the
+        Apple Neural Engine for x > ~10.4 (exp(x) exceeds 65504). The stable
+        decomposition shifts by the global max along the dim before exp.
+        """
+        if frontend == TorchFrontend.EXECUTORCH:
+            pytest.skip("torch._ops.aten.logcumsumexp.default is not Aten Canonical")
+
+        input_shape = (4, 6, 3, 5)
+
+        class LogCumSumExpModel(nn.Module):
+            def forward(self, x):
+                return torch.logcumsumexp(x, dim=axis)
+
+        model = LogCumSumExpModel().eval()
+        # Values chosen to overflow naive fp16: exp(15) ~ 3.3M, exp(20) ~ 485M.
+        # With the stable decomposition, exp(x - max(x)) is always in (0, 1].
+        input_values = torch.randn(*input_shape) * 15.0
+        self.run_compare_torch(
+            input_values,
+            model,
+            backend=backend,
+            compute_unit=compute_unit,
+            frontend=frontend,
+        )
+
 
 class TestHannWindow(TorchBaseTest):
     @pytest.mark.parametrize(