Skip to content

[torchlib] Fix torchvision::roi_align lowering to accept 7-arg schema#2830

Merged
justinchuby merged 5 commits intomicrosoft:mainfrom
FraGirla:fix-torch-roi-align
Mar 19, 2026
Merged

[torchlib] Fix torchvision::roi_align lowering to accept 7-arg schema#2830
justinchuby merged 5 commits intomicrosoft:mainfrom
FraGirla:fix-torch-roi-align

Conversation

@FraGirla
Copy link
Copy Markdown
Contributor

Dynamo ONNX export calls torchvision::roi_align with (input, rois, spatial_scale,
pooled_h, pooled_w, sampling_ratio, aligned), but torchlib’s lowering expected
output_size and rejected the extra args.

Update torchvision_roi_align to take pooled_h/pooled_w directly and add an OpInfo
input_wrangler to adapt wrapper-style roi_align(input, boxes, output_size, ...).

Refs pytorch/pytorch#175732
Test: pytest tests/function_libs/torch_lib/ops_test.py -k roi_align -v

@FraGirla
Copy link
Copy Markdown
Contributor Author

@microsoft-github-policy-service agree

Comment thread onnxscript/function_libs/torch_lib/ops/vision.py
@justinchuby justinchuby self-assigned this Feb 25, 2026
Comment thread tests/function_libs/torch_lib/ops_test_data.py
Comment thread tests/function_libs/torch_lib/ops_test_data.py Outdated
@codecov
Copy link
Copy Markdown

codecov Bot commented Feb 25, 2026

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 71.79%. Comparing base (3f197a2) to head (18cd621).
⚠️ Report is 10 commits behind head on main.

Additional details and impacted files
@@            Coverage Diff             @@
##             main    #2830      +/-   ##
==========================================
- Coverage   71.79%   71.79%   -0.01%     
==========================================
  Files         239      239              
  Lines       29018    29017       -1     
  Branches     2864     2864              
==========================================
- Hits        20833    20832       -1     
  Misses       7213     7213              
  Partials      972      972              

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

@justinchuby
Copy link
Copy Markdown
Collaborator

Still want to check when this was changed. Since it is a breaking change we need to define the function conditionally based on pytorch version.

@FraGirla
Copy link
Copy Markdown
Contributor Author

I checked versions: this fix targets the torch.export/Dynamo ONNX exporter path (torch.onnx.export(..., dynamo=True)), which is available in recent PyTorch releases (docs mention >= 2.5). The change simply aligns the torchlib lowering with the underlying operator schema for torchvision::roi_align (7 args: input, rois, spatial_scale, pooled_height, pooled_width, sampling_ratio, aligned).

The operator itself has been 7-arg for a long time; what’s new is that the Dynamo ONNX exporter path now hits the operator-level call directly, while the previous torchlib lowering matched the Python wrapper signature (output_size tuple), so the mismatch only surfaced once this path got exercised more.

Also fixed lintrunner issues in a follow-up commit.

Dynamo ONNX export calls torchvision::roi_align with (input, rois, spatial_scale,
pooled_h, pooled_w, sampling_ratio, aligned), but torchlib’s torchvision_roi_align
expected output_size and rejected the extra args.

Update torchvision_roi_align to take pooled_h/pooled_w directly and add an OpInfo
input_wrangler to adapt wrapper-style roi_align(input, boxes, output_size, ...).

Refs pytorch/pytorch#175732
Test: pytest tests/function_libs/torch_lib/ops_test.py -k roi_align -v
… remove redundant tests and simplify input handling
@FraGirla FraGirla force-pushed the fix-torch-roi-align branch from 3fd9cee to 18cd621 Compare February 27, 2026 08:45
@FraGirla
Copy link
Copy Markdown
Contributor Author

Also the Dynamo export is now the default in the torch 2.9 milestone. That's why the error came out right now
Source: pytorch/pytorch#151693

Copy link
Copy Markdown
Collaborator

@justinchuby justinchuby left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's just do it

@justinchuby justinchuby merged commit ebb0007 into microsoft:main Mar 19, 2026
32 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Development

Successfully merging this pull request may close these issues.

2 participants