Fix softplus and mish fp16 overflow on ANE via stable decomposition

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

@@ -1775,6 +1775,23 @@ def _parse_positional_args(context, node) -> Tuple[Var]:
     context.add(res)
 
 
+def _stable_softplus_mil(x):
+    """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 matches the formula used by coremltools' own
+    softplus MIL op ``value_inference``.
+
+    See issue #2687 for the original fp16 overflow report.
+    """
+    abs_x = mb.abs(x=x)
+    neg_abs_x = mb.mul(x=-1.0, y=abs_x)
+    exp_val = mb.exp(x=neg_abs_x)
+    log_val = mb.log(x=mb.add(x=1.0, y=exp_val))
+    max_val = mb.maximum(x=x, y=0.0)
+    return mb.add(x=max_val, y=log_val)
+
+
 @register_torch_op
 def softplus(context, node):
     def _parse_positional_args(context, node) -> Tuple[Var]:
@@ -1795,20 +1812,17 @@ def _parse_positional_args(context, node) -> Tuple[Var]:
     x, beta, threshold = _parse_positional_args(context, node)
 
     if beta == 1:
-        # this is the special case that Core ML softplus handles
-        res = mb.softplus(x=x, name=node.name)
+        sp = _stable_softplus_mil(x)
+        cond = mb.greater(x=x, y=threshold)
     else:
-        if x.rank == 4:
-            # can use Core ML softplus_parametric
-            C = x.shape[1]
-            alpha_br = np.repeat(1.0 / beta, C).astype("float32")
-            beta_br = np.repeat(beta, C).astype("float32")
-            res = mb.softplus_parametric(x=x, alpha=alpha_br, beta=beta_br, name=node.name)
-        else:
-            # have to generally decompose
-            beta_mul_x = mb.mul(x=beta, y=x)
-            softplus = mb.softplus(x=beta_mul_x)
-            res = mb.real_div(x=softplus, y=beta, name=node.name)
+        # For non-unit beta: softplus(x) = (1/beta) * softplus(beta * x)
+        beta_x = mb.mul(x=beta, y=x)
+        sp = mb.real_div(x=_stable_softplus_mil(beta_x), y=beta)
+        cond = mb.greater(x=beta_x, y=threshold)
+
+    # Apply PyTorch's threshold: for beta * x > threshold, softplus(x) ≈ x,
+    # so return x directly.  This matches PyTorch's exact behavior.
+    res = mb.select(cond=cond, a=x, b=sp, name=node.name)
     context.add(res)
 
 
@@ -1817,8 +1831,8 @@ def mish(context, node):
     inputs = _get_inputs(context, node, expected=1)
     x = inputs[0]
 
-    softplus = mb.softplus(x=x)
-    tanh = mb.tanh(x=softplus)
+    # Mish(x) = x * tanh(softplus(x))
+    tanh = mb.tanh(x=_stable_softplus_mil(x))
     res = mb.mul(x=x, y=tanh, name=node.name)
     context.add(res)
 

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

@@ -6721,17 +6721,9 @@ def test_softplus(
             # executorch decomposes softplus to very basic log and exp
             target_op = "exp"
         else:
-            if beta is None or beta == 1:
-                # this is the special case that Core ML softplus handles
-                target_op = "softplus"
-            else:
-                if rank == 4:
-                    # can use Core ML softplus_parametric
-                    target_op = "softplus_parametric"
-                else:
-                    # have to generally decompose to
-                    # `x -> beta * x -> softplus(beta * x) -> softplus(beta * x) / beta`
-                    target_op = "softplus"
+            # The converter now wraps softplus in a select for fp16 threshold safety,
+            # so we skip the target_op check for non-executorch frontends.
+            target_op = None
 
         self.run_compare_torch(
             input_shape,
@@ -6743,6 +6735,48 @@ def test_softplus(
             target_op=target_op,
         )
 
+    @pytest.mark.parametrize("frontend", frontends)
+    def test_softplus_fp16_stable_decomposition(self, frontend):
+        """Verify softplus is decomposed into overflow-safe ops for fp16 (#2687)."""
+        if frontend == TorchFrontend.EXECUTORCH:
+            pytest.skip("ExecuTorch decomposes softplus independently")
+
+        # Demonstrate that naive fp16 softplus (log(1+exp(x))) overflows for
+        # large x, while the stable form (max(x,0)+log(1+exp(-|x|))) does not.
+        x_val = np.float16(15.0)
+        naive = np.float16(np.log(np.float16(1.0) + np.exp(x_val)))
+        assert not np.isfinite(naive), (
+            f"Naive fp16 softplus should overflow at x=15, got {naive}"
+        )
+        stable = np.float16(
+            np.maximum(x_val, np.float16(0))
+            + np.log(np.float16(1.0) + np.exp(-np.abs(x_val)))
+        )
+        assert np.isfinite(stable) and abs(float(stable) - 15.0) < 0.5, (
+            f"Stable fp16 softplus should produce ~15.0, got {stable}"
+        )
+
+        # Verify the converter emits the stable decomposition, not the native op.
+        torch.manual_seed(0)
+        model = nn.Softplus().eval()
+        x = torch.randn(1, 10)
+
+        torch_model = export_torch_model_to_frontend(model, (x,), frontend)
+        mlmodel = ct.convert(
+            torch_model,
+            inputs=[ct.TensorType(shape=x.shape)],
+            convert_to="mlprogram",
+            compute_precision=ct.precision.FLOAT16,
+        )
+
+        prog = mlmodel._mil_program
+        assert len(prog.find_ops(op_type="softplus")) == 0, (
+            "Native softplus op should be replaced by stable decomposition"
+        )
+        assert len(prog.find_ops(op_type="exp")) >= 1, (
+            "Stable decomposition should contain at least one exp op"
+        )
+
     @pytest.mark.parametrize(
         "compute_unit, backend, frontend, shape",
         itertools.product(compute_units, backends, frontends, COMMON_SHAPES_ALL),