fix(tf-frontend): use numerically stable softplus decomposition#2756
Open
davidnichols-ops wants to merge 4 commits into
Open
fix(tf-frontend): use numerically stable softplus decomposition#2756davidnichols-ops wants to merge 4 commits into
davidnichols-ops wants to merge 4 commits into
Conversation
…e#2747) The TensorFlow frontend's Softplus op emitted the native mb.softplus op, which on the Apple Neural Engine is susceptible to fp16 overflow for x > ~10.4 (the log(1 + exp(x)) computation overflows in fp16). Replace mb.softplus with a shared stable_softplus_mil helper that decomposes softplus into max(x, 0) + log(1 + exp(-|x|)), matching the formula already used by coremltools' own softplus MIL op value_inference. Since -|x| <= 0, exp(-|x|) is always in (0, 1], so no overflow occurs in any precision. The Torch frontend already had a similar stable decomposition; this consolidates both frontends onto a shared helper. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Contributor
|
@TobyRoseman Could you review this numerically stable softplus fix when you have bandwidth? |
TobyRoseman
requested changes
Jul 15, 2026
| frontend_only=True, | ||
| ) | ||
|
|
||
| # res is a tuple: (spec, mlmodel, input_key_values, pred) |
Collaborator
There was a problem hiding this comment.
We typically don't directly inspect the MIL for unit tests like this. Please remove this part.
Collaborator
There was a problem hiding this comment.
You seem to have misunderstood what I was asking. I didn't want you to remove the whole unit test. I just wanted you to remove:
# res is a tuple: (spec, mlmodel, input_key_values, pred)
mlmodel = res[1]
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"
)- Inline the stable softplus decomposition directly into the Softplus op instead of a separate stable_softplus_mil helper in _utils.py - Remove test_softplus_fp16_stable_decomposition, which directly inspected the MIL program (not typical for unit tests here)
Keep the fp16 overflow demonstration and run_compare_tf conversion check, but drop the direct MIL program inspection per review.
Author
|
@TobyRoseman I've restored the test without the MIL inspection part (commit 5989825) — the fp16 overflow demonstration and standard |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Softplusop emitted the nativemb.softplusop, which on the Apple Neural Engine is susceptible to fp16 overflow forx > ~10.4(thelog(1 + exp(x))computation overflows in fp16). See issue TensorFlow frontend Softplus converter still emits native softplus op (fp16 overflow risk on ANE) #2747.stable_softplus_milhelper incoremltools/converters/mil/frontend/_utils.pythat decomposes softplus intomax(x, 0) + log(1 + exp(-|x|)), matching the formula already used by coremltools' own softplus MIL opvalue_inference. Since-|x| <= 0,exp(-|x|)is always in(0, 1], so no overflow occurs in any precision.Softplusop to use this helper instead ofmb.softplus.test_softplus_fp16_stable_decompositionthat verifies the converter emits the stable decomposition (no nativesoftplusop, containsexpandabsops).Deleted: nothing — the native
mb.softplusop call is replaced (not deleted from the codebase), and the stable decomposition is the only correct fix since the ANE's fp16 overflow is a hardware limitation that cannot be worked around at the op level.Test plan
stable_softplus_milhelper verified in isolation: builds correct MIL program with 0softplusops, 1expop, 1absoptest_softplus_fp16_stable_decompositiontest requires building coremltools native libraries (BlobWriter) — verified logic but could not run full pytest locallyGenerated with Devin