Add aten._trilinear support to torch_lib core#2850
Add aten._trilinear support to torch_lib core#2850WineChord wants to merge 6 commits intomicrosoft:mainfrom
Conversation
|
@microsoft-github-policy-service agree |
Codecov Report❌ Patch coverage is
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. |
There was a problem hiding this comment.
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::_trilinearlowering inonnxscript/function_libs/torch_lib/ops/core.pyby generating an einsum equation fromexpand*/sumdim. - Add OpInfo/sample-input coverage for
torch.ops.aten._trilinear.defaultintests/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. |
|
Can you consider the copilot comments as well? |
There was a problem hiding this comment.
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::_trilinearinonnxscript/function_libs/torch_lib/ops/core.py. - Add OpInfo/sample-input coverage for
ops.aten._trilinear.defaultand 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. |
| input_rank = getattr(i1, "rank", None) | ||
| if input_rank is None: | ||
| input_rank = len(i1.shape) | ||
| total_dim = input_rank + len(expand1) |
Summary
aten::_trilinearincore.pyexpand*/sumdimarguments as Python tuples during ONNX tracing tests_trilineardimension lists before building the einsum equationTesting
Context: requested from the review on pytorch/pytorch#177223.