[ONNX] Support complex in FX exporter #100554

titaiwangms · 2023-05-03T16:13:02Z

Stack from ghstack (oldest at bottom):

Previous to the PR, the complex dtype would only fail. This PR keeps torch.fx.Graph with complex dtype, while mapping them to float dtype in torchscript(onnx) graph with real representation.

The change happens in multiple files:

placeholder: Apply torch.view_as_real() before sending fake tensor to graph building.
call_function: Fill in TorchScriptTensor dtype and shape with real representation dtype and shape.
Registry: Add is_complex, and supports complex onnxfunction.
Dispatcher: Filter with/out complex onnxfunction before opschema matching, based on the dtype in torch args
Test cases: input/output view_as_real for result comparisons.

[ghstack-poisoned]

pytorch-bot · 2023-05-03T16:13:05Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/100554

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit ddb5acb:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

ghstack-source-id: 9758268225091a89c7dd1954bc5e3771ecb2f926 Pull Request resolved: #100554

torch/onnx/_internal/fx/op_validation.py

[ghstack-poisoned]

ghstack-source-id: cea8f8c8b39132d368439fd5b9ad208bcdeb0606 Pull Request resolved: #100554

torch/onnx/_internal/fx/passes/fx_to_onnxscript.py

torch/onnx/_type_utils.py

torch/onnx/_internal/fx/op_validation.py

justinchuby

Awesome!

torch/onnx/_type_utils.py

[ghstack-poisoned]

torch/onnx/_internal/fx/registration.py

torch/onnx/_internal/fx/onnxfunction_dispatcher.py

torch/onnx/_internal/fx/type_utils.py

justinchuby · 2023-07-26T23:58:57Z

LGTM. Just a few minor comments

Previous to the PR, the complex dtype would only fail. This PR keeps torch.fx.Graph with complex dtype, while mapping them to float dtype in torchscript(onnx) graph with real representation. The change happens in multiple files: 1. `placeholder`: Apply torch.view_as_real() before sending fake tensor to graph building. 2. `call_function`: Fill in TorchScriptTensor dtype and shape with real representation dtype and shape. 3. Registry: Add `is_complex`, and supports complex onnxfunction. 4. Dispatcher: Filter with/out complex onnxfunction before opschema matching, based on the dtype in torch args 5. Test cases: input/output view_as_real for result comparisons. [ghstack-poisoned]

ghstack-source-id: a722217c4578f73ede800118daeb1932f74dbf06 Pull Request resolved: #100554

Previous to the PR, the complex dtype would only fail. This PR keeps torch.fx.Graph with complex dtype, while mapping them to float dtype in torchscript(onnx) graph with real representation. The change happens in multiple files: 1. `placeholder`: Apply torch.view_as_real() before sending fake tensor to graph building. 2. `call_function`: Fill in TorchScriptTensor dtype and shape with real representation dtype and shape. 3. Registry: Add `is_complex`, and supports complex onnxfunction. 4. Dispatcher: Filter with/out complex onnxfunction before opschema matching, based on the dtype in torch args 5. Test cases: input/output view_as_real for result comparisons. [ghstack-poisoned]

BowenBao · 2023-07-27T16:58:10Z

test/onnx/onnx_test_common.py

+        # NOTE: ONNX Runtime doesn't support tensor of complex64/complex128, so we
+        # convert them to float32/float64.
+        # TODO: Need stft enabled to test complex64/complex128
+        if isinstance(ref_output, torch.Tensor) and torch.is_complex(ref_output):


You should add a new input/output adaptor step for this. Checkout RemoveNonTensorInputStep for example.
Feel free to do it in this or follow-up PR.

BowenBao · 2023-07-27T16:58:34Z

test/onnx/onnx_test_common.py

+        else input
+        for input in pytorch_inputs
+    ]
+    ort_input = {k: v.cpu().numpy() for k, v in zip(input_names, pytorch_inputs)}


ditto: input adapter step

BowenBao

LGTM w/ comments. Please also resolve others' comments before landing.

Previous to the PR, the complex dtype would only fail. This PR keeps torch.fx.Graph with complex dtype, while mapping them to float dtype in torchscript(onnx) graph with real representation. The change happens in multiple files: 1. `placeholder`: Apply torch.view_as_real() before sending fake tensor to graph building. 2. `call_function`: Fill in TorchScriptTensor dtype and shape with real representation dtype and shape. 3. Registry: Add `is_complex`, and supports complex onnxfunction. 4. Dispatcher: Filter with/out complex onnxfunction before opschema matching, based on the dtype in torch args 5. Test cases: input/output view_as_real for result comparisons. [ghstack-poisoned]

justinchuby · 2023-07-27T19:23:47Z

Are we handling any complex constants that may be inside the graph? Maybe that’s an unlikely case that we don’t need to care yet?

Previous to the PR, the complex dtype would only fail. This PR keeps torch.fx.Graph with complex dtype, while mapping them to float dtype in torchscript(onnx) graph with real representation. The change happens in multiple files: 1. `placeholder`: Apply torch.view_as_real() before sending fake tensor to graph building. 2. `call_function`: Fill in TorchScriptTensor dtype and shape with real representation dtype and shape. 3. Registry: Add `is_complex`, and supports complex onnxfunction. 4. Dispatcher: Filter with/out complex onnxfunction before opschema matching, based on the dtype in torch args 5. Test cases: input/output view_as_real for result comparisons. [ghstack-poisoned]

titaiwangms · 2023-07-27T21:25:54Z

Are we handling any complex constants that may be inside the graph? Maybe that’s an unlikely case that we don’t need to care yet?

It's not covered in this PR. We can add it when we have a model or use case.

Previous to the PR, the complex dtype would only fail. This PR keeps torch.fx.Graph with complex dtype, while mapping them to float dtype in torchscript(onnx) graph with real representation. The change happens in multiple files: 1. `placeholder`: Apply torch.view_as_real() before sending fake tensor to graph building. 2. `call_function`: Fill in TorchScriptTensor dtype and shape with real representation dtype and shape. 3. Registry: Add `is_complex`, and supports complex onnxfunction. 4. Dispatcher: Filter with/out complex onnxfunction before opschema matching, based on the dtype in torch args 5. Test cases: input/output view_as_real for result comparisons. [ghstack-poisoned]

titaiwangms · 2023-07-27T23:42:44Z

@pytorchbot merge

pytorchmergebot · 2023-07-27T23:44:38Z

Merge started

Your 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

Advanced Debugging

Check the merge workflow status
here

pytorchmergebot · 2023-07-28T05:43:11Z

The merge job was canceled. If you believe this is a mistake,then you can re trigger it through pytorch-bot.

titaiwangms · 2023-07-28T07:01:18Z

@pytorchbot merge

pytorchmergebot · 2023-07-28T07:03:01Z

Merge started

Your 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

Advanced Debugging

Check the merge workflow status
here

Previous to the PR, the complex dtype would only fail. This PR keeps torch.fx.Graph with complex dtype, while mapping them to float dtype in torchscript(onnx) graph with real representation. The change happens in multiple files: 1. `placeholder`: Apply torch.view_as_real() before sending fake tensor to graph building. 2. `call_function`: Fill in TorchScriptTensor dtype and shape with real representation dtype and shape. 3. Registry: Add `is_complex`, and supports complex onnxfunction. 4. Dispatcher: Filter with/out complex onnxfunction before opschema matching, based on the dtype in torch args 5. Test cases: input/output view_as_real for result comparisons. Pull Request resolved: pytorch#100554 Approved by: https://github.com/BowenBao

[ONNX] Support complex in FX exporter

67b3a7c

[ghstack-poisoned]

titaiwangms requested review from BowenBao and abock as code owners May 3, 2023 16:13

pytorch-bot bot added the release notes: onnx torch.onnx related changes that should show up in the release notes label May 3, 2023

titaiwangms added a commit that referenced this pull request May 3, 2023

[ONNX] Support complex in FX exporter

e38fbe3

ghstack-source-id: 9758268225091a89c7dd1954bc5e3771ecb2f926 Pull Request resolved: #100554

titaiwangms marked this pull request as draft May 3, 2023 16:13

titaiwangms commented May 3, 2023

View reviewed changes

torch/onnx/_internal/fx/op_validation.py Outdated Show resolved Hide resolved

pytorchbot added the open source label May 3, 2023

titaiwangms added module: onnx Related to torch.onnx topic: new features topic category labels Jun 18, 2023

Update on "[ONNX] Support complex in FX exporter"

53aaf9b

[ghstack-poisoned]

titaiwangms marked this pull request as ready for review June 19, 2023 03:44

titaiwangms requested a review from justinchuby June 19, 2023 03:57

Update on "[ONNX] Support complex in FX exporter"

85a5246

[ghstack-poisoned]

titaiwangms added a commit that referenced this pull request Jun 19, 2023

[ONNX] Support complex in FX exporter

901ef00

ghstack-source-id: cea8f8c8b39132d368439fd5b9ad208bcdeb0606 Pull Request resolved: #100554

justinchuby self-assigned this Jun 19, 2023

justinchuby reviewed Jun 19, 2023

View reviewed changes

torch/onnx/_internal/fx/passes/fx_to_onnxscript.py Outdated Show resolved Hide resolved

justinchuby reviewed Jun 19, 2023

View reviewed changes

torch/onnx/_type_utils.py Outdated Show resolved Hide resolved

justinchuby reviewed Jun 19, 2023

View reviewed changes

torch/onnx/_internal/fx/op_validation.py Outdated Show resolved Hide resolved

justinchuby previously approved these changes Jun 19, 2023

View reviewed changes

thiagocrepaldi reviewed Jun 19, 2023

View reviewed changes

torch/onnx/_type_utils.py Outdated Show resolved Hide resolved

titaiwangms mentioned this pull request Jun 21, 2023

[ONNX] Separate fx _type_utils from torchscript exporter #103942

Closed

Update on "[ONNX] Support complex in FX exporter"

18c827c

[ghstack-poisoned]

titaiwangms mentioned this pull request Jun 21, 2023

[ONNX] Refactor FX Registry and Support Custom Operator in FX Exporter #103943

Closed

Update on "[ONNX] Support complex in FX exporter"

52ddbe5

[ghstack-poisoned]

Update on "[ONNX] Support complex in FX exporter"

24d5d7d

[ghstack-poisoned]

justinchuby reviewed Jul 26, 2023

View reviewed changes

torch/onnx/_internal/fx/registration.py Outdated Show resolved Hide resolved

titaiwangms commented Jul 26, 2023

View reviewed changes

torch/onnx/_internal/fx/onnxfunction_dispatcher.py Outdated Show resolved Hide resolved

justinchuby reviewed Jul 26, 2023

View reviewed changes

torch/onnx/_internal/fx/type_utils.py Outdated Show resolved Hide resolved

torch/onnx/_internal/fx/type_utils.py Outdated Show resolved Hide resolved

titaiwangms added a commit that referenced this pull request Jul 27, 2023

[ONNX] Support complex in FX exporter

6c9e034

ghstack-source-id: a722217c4578f73ede800118daeb1932f74dbf06 Pull Request resolved: #100554

titaiwangms mentioned this pull request Jul 27, 2023

[ONNX] Expose OnnxRegistry publicly #106140

Closed

BowenBao reviewed Jul 27, 2023

View reviewed changes

BowenBao approved these changes Jul 27, 2023

View reviewed changes

titaiwangms added the ciflow/trunk Trigger trunk jobs on your pull request label Jul 27, 2023

pytorchmergebot added the merging label Jul 27, 2023

pytorchmergebot added Merged and removed merging labels Jul 28, 2023

pytorchmergebot closed this in eebfb92 Jul 28, 2023

facebook-github-bot deleted the gh/titaiwangms/17/head branch July 31, 2023 14:17

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[ONNX] Support complex in FX exporter #100554

[ONNX] Support complex in FX exporter #100554

titaiwangms commented May 3, 2023 •

edited

Loading

pytorch-bot bot commented May 3, 2023 •

edited

Loading

justinchuby left a comment

justinchuby commented Jul 26, 2023

BowenBao Jul 27, 2023

titaiwangms Jul 27, 2023

BowenBao Jul 27, 2023

titaiwangms Jul 27, 2023

BowenBao left a comment

justinchuby commented Jul 27, 2023

titaiwangms commented Jul 27, 2023

titaiwangms commented Jul 27, 2023

pytorchmergebot commented Jul 27, 2023

pytorchmergebot commented Jul 28, 2023

titaiwangms commented Jul 28, 2023

pytorchmergebot commented Jul 28, 2023

[ONNX] Support complex in FX exporter #100554

[ONNX] Support complex in FX exporter #100554

Conversation

titaiwangms commented May 3, 2023 • edited Loading

pytorch-bot bot commented May 3, 2023 • edited Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/100554

✅ No Failures

justinchuby left a comment

Choose a reason for hiding this comment

justinchuby commented Jul 26, 2023

BowenBao Jul 27, 2023

Choose a reason for hiding this comment

titaiwangms Jul 27, 2023

Choose a reason for hiding this comment

BowenBao Jul 27, 2023

Choose a reason for hiding this comment

titaiwangms Jul 27, 2023

Choose a reason for hiding this comment

BowenBao left a comment

Choose a reason for hiding this comment

justinchuby commented Jul 27, 2023

titaiwangms commented Jul 27, 2023

titaiwangms commented Jul 27, 2023

pytorchmergebot commented Jul 27, 2023

Merge started

pytorchmergebot commented Jul 28, 2023

titaiwangms commented Jul 28, 2023

pytorchmergebot commented Jul 28, 2023

Merge started

titaiwangms commented May 3, 2023 •

edited

Loading

pytorch-bot bot commented May 3, 2023 •

edited

Loading