Skip to content

Add aten._trilinear support to torch_lib core#2850

Open
WineChord wants to merge 6 commits intomicrosoft:mainfrom
WineChord:add_trilinear_support
Open

Add aten._trilinear support to torch_lib core#2850
WineChord wants to merge 6 commits intomicrosoft:mainfrom
WineChord:add_trilinear_support

Conversation

@WineChord
Copy link

@WineChord WineChord commented Mar 12, 2026

Summary

  • add an opset 18 torch_lib lowering for aten::_trilinear in core.py
  • add OpInfo coverage and keep expand*/sumdim arguments as Python tuples during ONNX tracing tests
  • validate _trilinear dimension lists before building the einsum equation

Testing

  • env -u LD_LIBRARY_PATH pytest tests/function_libs/torch_lib/ops_test.py -k "_trilinear" -vv
  • env -u LD_LIBRARY_PATH pytest tests/function_libs/torch_lib/ops_test.py -k "bilinear and not upsample" -vv

Context: requested from the review on pytorch/pytorch#177223.

@WineChord
Copy link
Author

@microsoft-github-policy-service agree

@codecov
Copy link

codecov bot commented Mar 12, 2026

Codecov Report

❌ Patch coverage is 75.00000% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 71.87%. Comparing base (4c4f7a0) to head (c5d4cda).

Files with missing lines Patch % Lines
onnxscript/function_libs/torch_lib/ops/core.py 75.00% 4 Missing and 4 partials ⚠️
Additional details and impacted files
@@           Coverage Diff           @@
##             main    #2850   +/-   ##
=======================================
  Coverage   71.86%   71.87%           
=======================================
  Files         239      239           
  Lines       29139    29171   +32     
  Branches     2875     2880    +5     
=======================================
+ Hits        20942    20966   +24     
- Misses       7219     7223    +4     
- Partials      978      982    +4     

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

Copy link
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.

Thanks! Minor nits

@justinchuby justinchuby added the module: torchlib Related to the torch/aten function lib in development label Mar 12, 2026
Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull request overview

Adds a torch_lib lowering for aten::_trilinear (via ONNX Einsum) and wires it into the torch-lib OpInfo-based test harness to exercise the new lowering.

Changes:

  • Implement aten::_trilinear lowering in onnxscript/function_libs/torch_lib/ops/core.py by generating an einsum equation from expand*/sumdim.
  • Add OpInfo/sample-input coverage for torch.ops.aten._trilinear.default in tests/function_libs/torch_lib/extra_opinfo.py.
  • Register the new op test entry in tests/function_libs/torch_lib/ops_test_data.py.

Reviewed changes

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

File Description
tests/function_libs/torch_lib/ops_test_data.py Registers _trilinear in the torch_lib op test list.
tests/function_libs/torch_lib/extra_opinfo.py Adds OpInfo + sample inputs to drive _trilinear tests.
onnxscript/function_libs/torch_lib/ops/core.py Implements aten::_trilinear lowering using ONNX Einsum.

@justinchuby
Copy link
Collaborator

Can you consider the copilot comments as well?

Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull request overview

This PR adds TorchLib support for PyTorch’s internal aten::_trilinear op by lowering it to ONNX Einsum, and extends the TorchLib OpInfo test matrix to cover the new op.

Changes:

  • Add an opset18 TorchLib lowering for aten::_trilinear in onnxscript/function_libs/torch_lib/ops/core.py.
  • Add OpInfo/sample-input coverage for ops.aten._trilinear.default and register it in TorchLib test data.
  • Add dimension-list validation helpers for building the Einsum equation.

Reviewed changes

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

File Description
tests/function_libs/torch_lib/ops_test_data.py Registers ops.aten._trilinear.default in TorchLib operator test data.
tests/function_libs/torch_lib/extra_opinfo.py Adds OpInfo + sample inputs for _trilinear to exercise tracing/export and numerical checks.
onnxscript/function_libs/torch_lib/ops/core.py Implements _trilinear lowering via op.Einsum and adds helpers to build/validate the equation.

Comment on lines +9863 to +9866
input_rank = getattr(i1, "rank", None)
if input_rank is None:
input_rank = len(i1.shape)
total_dim = input_rank + len(expand1)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

module: torchlib Related to the torch/aten function lib in development

Projects

Development

Successfully merging this pull request may close these issues.

4 participants