Fix attention mask to use float_lowest instead of -inf and add NaN-safe softmax handling

[Aniketsy]

#2561

• Use lowest representable float value instead of -inf for attention masks.
• Add NaN-safe handling and a unit test for softmax with all masked positions.

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 !

[Aniketsy]

@microsoft-github-policy-service agree

[justinchuby]

If this PR is facilitated by an AI, please disclose its usage.

[codecov]

Codecov Report

❌ Patch coverage is 0% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 70.45%. Comparing base (a9cb429) to head (72b2c2c).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
onnxscript/function_libs/torch_lib/ops/nn.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2654   +/-   ##
=======================================
  Coverage   70.45%   70.45%           
=======================================
  Files         228      228           
  Lines       27177    27177           
  Branches     2734     2734           
=======================================
  Hits        19148    19148           
  Misses       7092     7092           
  Partials      937      937           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[Aniketsy]

I went through the questions you mentioned and yes, i used AI assistance to help add the unit test.

[justinchuby]

I think this is reasonable. @xadupre @titaiwangms what do you think?

[titaiwangms]

Yeah, that's also what ort folk would suggest: microsoft/onnxruntime#27129 (comment)

[justinchuby]

@Aniketsy could you fix the lint issues? Thanks

[Aniketsy]

@justinchuby fixed lint error.