Fix aten_unbind for torch >= 2.7 dynamo export

[afshin-paydar]

Fix aten_unbind for torch >= 2.7 dynamo export

Problem

When exporting PyTorch models to ONNX with dynamo=True on torch >= 2.7, the aten_unbind operation fails with:

TypeError: 'SymbolicTensor' object is not iterable

This occurs because op.Split(self, axis=dim, num_outputs=num_outputs) returns a single SymbolicTensor object rather than an iterable sequence during dynamo export. The subsequent list comprehension [op.Squeeze(out, [dim]) for out in outputs] attempts to iterate over this non-iterable object, causing that error.

This affects models using LSTM and other operations that internally call unbind, preventing successful ONNX export.

Solution

Replace the Split op approach with explicit Slice operations:

• For each output along the unbind dimension, create an individual Slice operation to extract one element
• Apply Squeeze to remove the size-1 dimension
• Collect all results in a list

Test coverage:

• num_outputs = 1: test_unbind_size_one
• num_outputs = 2: test_unbind_with_lstm, test_unbind_dim1 (size 2 along some dims)
• num_outputs = 3: test_unbind_dim0, test_unbind_dim1
• num_outputs = 4: test_unbind_negative_dim
• Negative dimensions: test_unbind_negative_dim

Related Issues

Fixes pytorch/pytorch#168969

[afshin-paydar]

@microsoft-github-policy-service agree

[codecov]

Codecov Report

❌ Patch coverage is 85.71429% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 70.11%. Comparing base (9dbf685) to head (9825c36).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
onnxscript/function_libs/torch_lib/ops/core.py	85.71%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2719   +/-   ##
=======================================
  Coverage   70.10%   70.11%           
=======================================
  Files         226      226           
  Lines       27226    27228    +2     
  Branches     2747     2747           
=======================================
+ Hits        19086    19090    +4     
+ Misses       7194     7193    -1     
+ Partials      946      945    -1     

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

[titaiwangms]

Can you add your test as well? https://github.com/microsoft/onnxscript/blob/main/tests/function_libs/torch_lib/e2e_ops_tests.py

[justinchuby]

Thanks. The fix is reasonable (and arguably better than the original because it's easier to optimize). Although I don't know why the exporter doesn't handle split correctly (it should)