diff --git a/coremltools/converters/mil/frontend/tensorflow/ops.py b/coremltools/converters/mil/frontend/tensorflow/ops.py index 1f45e821e..ebb07e0ff 100644 --- a/coremltools/converters/mil/frontend/tensorflow/ops.py +++ b/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) diff --git a/coremltools/converters/mil/frontend/tensorflow/test/test_ops.py b/coremltools/converters/mil/frontend/tensorflow/test/test_ops.py index af78a8475..4109c8483 100644 --- a/coremltools/converters/mil/frontend/tensorflow/test/test_ops.py +++ b/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( diff --git a/coremltools/converters/mil/frontend/torch/ops.py b/coremltools/converters/mil/frontend/torch/ops.py index eea3a50d0..0abde2d20 100644 --- a/coremltools/converters/mil/frontend/torch/ops.py +++ b/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) diff --git a/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py b/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py index 8f9f640a8..10d10c30d 100644 --- a/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py +++ b/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(