Add constant folding of Split with all constant inputs

[iksnagreb]

Adds constant folding of Split with all constant inputs as well as replacing single-output Splits by Identity via the registry. Also adds 5 simple test cases covering both options: Second input specifying the sizes per split as well as nearly even splits into num_outputs via the attribute.

[Copilot]

Pull Request Overview

This PR adds constant folding optimization for the Split operator when all inputs are constants, replacing them with individual Constant nodes. It also optimizes single-output Split operations by replacing them with Identity nodes.

• Implements constant folding for Split operations with all constant inputs
• Adds optimization to replace single-output Splits with Identity operations
• Includes comprehensive test coverage for both num_outputs and splits parameter variants

Reviewed Changes

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

File	Description
onnxscript/optimizer/_constant_folding.py	Adds the split function to handle constant folding and Identity replacement for Split operations
onnxscript/optimizer/_constant_folding_test.py	Adds 5 test cases covering Identity replacement and constant folding scenarios

[justinchuby]

Thank you!

[justinchuby]

LGTM. @gramalingam what do you think?

[codecov]

Codecov Report

❌ Patch coverage is 94.44444% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.13%. Comparing base (e8005e9) to head (89d57a0).
⚠️ Report is 11 commits behind head on main.

Files with missing lines	Patch %	Lines
onnxscript/optimizer/_constant_folding.py	84.00%	2 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2506      +/-   ##
==========================================
+ Coverage   70.08%   70.13%   +0.05%     
==========================================
  Files         215      215              
  Lines       25810    25938     +128     
  Branches     2587     2612      +25     
==========================================
+ Hits        18089    18192     +103     
- Misses       6823     6841      +18     
- Partials      898      905       +7     

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

[gramalingam]

I think a better, more general, solution would be to fix the restriction in the general constant-folding (this else branch) ... so the underlying evaluator evaluates the Split op ... and we just need to bind each output value to the corresponding constant value.

[justinchuby]

Could you share what needs to be done to enable that part?

[justinchuby]

@iksnagreb as suggested here, we will instead fix handling for multi-output nodes. Please feel free to contribute or provide your inputs, thanks!

[gramalingam]

Good question. What needs to be done is create an appropriate Replacement object which wraps the list of new nodes (a list of Constant nodes like in current PR), and a list of values (which should be more or less what is returned by the reference implementation).

I realize that this is lacking even in the current PR down below ... if we don't specify the list of new values (to replace the list of old values produced by the constant-folded node), we are not going to connect these values to the consumers of the old values.

So, if we added a test-case where the outputs of Split are used subsequently (or if we test the graph outputs in the tests below actually refer to the correct split const values), we will likely find that this PR doesn't work. It stems from a limitation in the core constant-folding logic that needs to be fixed.

[justinchuby]

Blocking for now according to reviews, thanks!