Update constant folding behavior for large tensors

[justinchuby]

Suggested by #2466, I updated the constant folder logic to allow

Constant folding customization:

• Replaced the always_fold_ops parameter with a should_fold callable that determines on a per-node basis whether folding should occur. This allows users to specify more complex folding policies and makes the API more explicit. (FoldConstantsPass, fold_constants) [1] [2] [3] [4] [5]

Logging and diagnostics improvements:

• Upgraded logging throughout the folding process to provide more informative messages, including reasons for skipping nodes (e.g., control flow, non-deterministic ops, large inputs, or graph inputs) and explicit logging when should_fold returns a decision. [1] [2] [3]

Code cleanup and minor fixes:

• Removed the unused _update_type function.

Fix #2466

cc @iksnagreb

[codecov]

Codecov Report

❌ Patch coverage is 76.31579% with 9 lines in your changes missing coverage. Please review.
✅ Project coverage is 69.94%. Comparing base (2cc2502) to head (b99fe88).
⚠️ Report is 7 commits behind head on main.
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
onnxscript/optimizer/_constant_folding.py	67.85%	5 Missing and 4 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2488      +/-   ##
==========================================
- Coverage   70.00%   69.94%   -0.06%     
==========================================
  Files         215      216       +1     
  Lines       25988    26064      +76     
  Branches     2606     2614       +8     
==========================================
+ Hits        18192    18231      +39     
- Misses       6896     6932      +36     
- Partials      900      901       +1     

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

[justinchuby]

I will remove the allow bloat option for now. Will create a function arg for user specified rules

[justinchuby]

@iksnagreb @gramalingam I removed the previous always_fold option and added a should_fold arg that takes a function. This is a breaking change. Please let me know if it looks ok.

[justinchuby]

@gramalingam could you reapprove? Thanks