Always fold the Transpose node in the constant folder

[justinchuby]

• Create an always_fold_ops option to allow users to specify which ops should always be folded
• Refactored the FoldConstantsPass to hide internal attributes
• Update logic to check for graph initialized inputs and removed the need for tracking in the object states

[Copilot]

Pull Request Overview

Adds support for always-folding specified ops (defaulting to Transpose) in the constant‐folder pass and refactors private helper naming and state tracking.

• Introduce an always_fold_ops parameter to FoldConstantsPass so users can always fold selected op types.
• Rename and scope helper functions/attributes (e.g. prefixing with _) for clarity.
• Replace custom initializer‐input tracking and input‐size gating logic with a unified const‐value check and debug logging.

Comments suppressed due to low confidence (4)

onnxscript/optimizer/_constant_folding.py:1108

• Removed push_initializer_inputs/pop_initializer_inputs logic may change constant‐folding behavior for default‐valued inputs; verify this was intentional or reintroduce initializer tracking.

for node in graph:

onnxscript/optimizer/_constant_folding.py:859

• No tests cover the new always_fold_ops feature; consider adding unit tests to ensure ops like Transpose are always folded regardless of input size.

always_fold_ops: Collection[str] = frozenset(["Transpose"]),

onnxscript/optimizer/_constant_folding.py:844

• [nitpick] The class doc briefly lists attributes but it may help future readers to add parameter descriptions for always_fold_ops, especially its default behavior, to the __init__ docstring.

class FoldConstantsPass(ir.passes.InPlacePass):

onnxscript/optimizer/_constant_folding.py:864

• [nitpick] The parameter is typed as Collection[str] but the attribute uses frozenset[str]; consider using FrozenSet[str] or Collection[str] consistently to avoid confusion.

self.always_fold_ops: frozenset[str] = frozenset(always_fold_ops)

[codecov]

Codecov Report

Attention: Patch coverage is 89.55224% with 7 lines in your changes missing coverage. Please review.

Project coverage is 70.21%. Comparing base (325b2de) to head (3cacd86).
Report is 3 commits behind head on main.

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

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2355   +/-   ##
=======================================
  Coverage   70.20%   70.21%           
=======================================
  Files         198      198           
  Lines       24871    24877    +6     
  Branches     2659     2661    +2     
=======================================
+ Hits        17461    17467    +6     
+ Misses       6496     6491    -5     
- Partials      914      919    +5     

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

[Copilot]

Pull Request Overview

This pull request adds an option (always_fold_ops) to ensure that specific ops (e.g., Transpose) are always folded by the constant folding pass. It also refactors the FoldConstantsPass to use public attribute names and updates tests to reflect folding behavior under different input size limits.

• Renames and updates tests to verify that input size limits correctly govern folding behavior for Mul and Transpose ops.
• Refactors internal attributes in FoldConstantsPass and updates related logging and API calls.
• Introduces the always_fold_ops option across the constant folding API.

Reviewed Changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
onnxscript/optimizer/_constant_folding_test.py	Renamed and updated tests for input size limits and ensured Transpose op is always folded.
onnxscript/optimizer/_constant_folding.py	Refactored internal state attributes, updated utility function calls, and added always_fold_ops support.

Comments suppressed due to low confidence (2)

onnxscript/optimizer/_constant_folding_test.py:538

• [nitpick] The test function was renamed to 'test_input_size_limit' and now covers both Mul and Transpose folding behavior. Consider splitting this into separate tests or renaming the function to clearly reflect that it tests different ops for clarity and maintainability.

def test_large_transpose(self):

onnxscript/optimizer/_constant_folding.py:859

• [nitpick] The use of frozenset for always_fold_ops is acceptable, but please ensure that exposing this as a public parameter does not conflict with future API extensions or lead to unintentional behavior changes.

always_fold_ops: Collection[str] = frozenset(["Transpose"]),

[Copilot]

Pull Request Overview

This PR introduces an option to always fold specified ops—most notably the Transpose op—in the constant folding pass, refactors internal state management, and updates relevant tests.

• Introduces an always_fold_ops parameter to dictate ops that must always be folded regardless of tensor size limits.
• Refactors internal attributes in FoldConstantsPass to use internal (underscore-prefixed) names and updates helper functions accordingly.
• Adjusts test cases to verify behavior changes for input size limits and ensures Transpose is consistently folded.

Reviewed Changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
onnxscript/optimizer/_constant_folding_test.py	Updated tests to reflect new input size limits and the always-fold Transpose behavior.
onnxscript/optimizer/_constant_folding.py	Refactored internal state management, updated op-checking functions, and added parameter support for always folding ops.

Comments suppressed due to low confidence (1)

onnxscript/optimizer/_constant_folding.py:846

• Consider adding documentation on the expected format for op names in 'always_fold_ops' (for example, whether the op should include a domain prefix or not) to help future maintainers understand and use this parameter.

def __init__(..., always_fold_ops: Collection[str] = frozenset(["Transpose"]),) -> None:

[justinchuby]

@gramalingam