Fix attention mask to use float_lowest instead of -inf and add NaN-safe softmax handling #2654
Conversation
…est for softmax NaN case
|
@microsoft-github-policy-service agree |
| ) | ||
| raise | ||
|
|
||
| def test_softmax_with_all_inf_mask(): |
There was a problem hiding this comment.
Please remove the test. It does not belong here
| # Turn the Boolean mask to float: attn_mask.masked_fill(not attn_mask, -float('inf')) | ||
| zero = op.Constant(value=ir.tensor(0.0, dtype=query.dtype)) | ||
| neg_inf = op.Constant(value=ir.tensor(-float("inf"), dtype=query.dtype)) | ||
| neg_inf = op.Constant(value=ir.tensor(float_lowest(query.dtype)), dtype=query.dtype) |
There was a problem hiding this comment.
Actually, you can use query.dtype.min directly because it is an implemented method in ir.DataType: https://onnx.ai/ir-py/api/generated/onnx_ir.DataType.html#onnx_ir.DataType.min
|
If this PR is facilitated by an AI, please disclose its usage. |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| # Turn the Boolean mask to float: attn_mask.masked_fill(not attn_mask, -float('inf')) | ||
| zero = op.Constant(value=ir.tensor(0.0, dtype=query.dtype)) | ||
| neg_inf = op.Constant(value=ir.tensor(-float("inf"), dtype=query.dtype)) | ||
| neg_inf = op.Constant(value=ir.tensor(float_lowest(query.dtype)), dtype=query.dtype) |
Check failure
Code scanning / lintrunner
PYLINT/E1123 Error
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2654 +/- ##
==========================================
- Coverage 70.46% 70.46% -0.01%
==========================================
Files 224 224
Lines 26572 26575 +3
Branches 2637 2637
==========================================
+ Hits 18723 18725 +2
- Misses 6928 6929 +1
Partials 921 921 ☔ View full report in Codecov by Sentry. |
|
I went through the questions you mentioned and yes, i used AI assistance to help add the unit test. |
2 participants
#2561
Please let me know if my approach or fix needs any improvements . I’m open to feedback and happy to make changes based on suggestions.
Thankyou !