[ONNX] Support aten::unflatten in torchscript exporter #99056
Conversation
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/99056
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit fa9dae0: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
| @@ -462,6 +463,50 @@ def reduce_dim(g, self, dim, keepdim, dtype): | |||
| return reduce | |||
|
|
|||
|
|
|||
| @_onnx_symbolic("aten::unflatten") | |||
| @_beartype.beartype | |||
| def unflatten(g: jit_utils.GraphContext, input, dim, unflattened_size): | |||
There was a problem hiding this comment.
Comment why this requires minimum opset 13?
There was a problem hiding this comment.
I tried to put this in Opset9, and turns out ops like Concat, Slice, .., has moved attrs to inputs, and that complicates things in a way that we have to engage helper functions. By demand, opset 13 should be enough for current requests. I will put a comment here.
| @@ -462,6 +463,50 @@ def reduce_dim(g, self, dim, keepdim, dtype): | |||
| return reduce | |||
|
|
|||
|
|
|||
| @_onnx_symbolic("aten::unflatten") | |||
| @_beartype.beartype | |||
| def unflatten(g: jit_utils.GraphContext, input, dim, unflattened_size): | |||
There was a problem hiding this comment.
@justinchuby @titaiwangms Should we mention something like "implementation ported from torchlib + link"?
There was a problem hiding this comment.
I’m fine either way. Any considerations?
There was a problem hiding this comment.
Just a random thought, mostly for awareness. But it could get outdated 🤷
Maybe adding like "... originally ported ..." reduces the harm of getting outdated, and surfaces that.
There was a problem hiding this comment.
A comment referencing a link to onnx Function?
There was a problem hiding this comment.
Added a comment with link and function name to point out the source.
test/onnx/test_op_consistency.py
Outdated
| @@ -330,6 +332,7 @@ def reason_flaky() -> str: | |||
| fixme("ceil", dtypes=[torch.float64], reason=reason_onnx_runtime_does_not_support("Ceil", ["f64"])), | |||
| dont_care("sqrt", dtypes=BOOL_TYPES, reason=reason_onnx_does_not_support("Sqrt")), | |||
| dont_care("stft", opsets=[opsets_before(17)], reason=reason_onnx_does_not_support("STFT")), | |||
| dont_care("unflatten", opsets=[opsets_before(13)], reason=reason_onnx_does_not_support("Unflatten")), | |||
There was a problem hiding this comment.
Is the reason accurate? I don’t think there is an Unflatten op in onnx?
There was a problem hiding this comment.
Updated with fixme. The unsupport reason is only that using the helper function would have to rewrite the structure. fixme should be sufficient as a starting point if one needs the op in legacy opset version.
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Stack from ghstack (oldest at bottom):
Fixes #98857
Fixes #98190