[ONNX] Adds Support for Some Bitwise Ops in Onnx Exporter

[zjgarvey]

Addresses #126194

Adds support for

• "aten::bitwise_right_shift"
• "aten::bitwise_left_shift"
• "aten::bitwise_and"

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 508aeb7 with merge base 82c66bc ():
💚 Looks good so far! There are no failures yet. 💚

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

[zjgarvey]

To be completely honest, I'm not very familiar with the internals of torch.onnx.export.

Are decomposition passes applied to a graph before converting torch fx nodes to onnx nodes? I also believe there is some automatic dtype promotion that occurs, but I got errors when leaving out the explicit type promotion in the conversion for "aten::bitwise_and".

In any case, let me know if there is a better way of adding support for these ops and I can make changes.

[zjgarvey]

I install the official torch==2.3.0 using pip , the number of lines in my code for symbolic_opset18.py is much less than the one in current main branch,( only defining col2im) and then I added the changes from your two code files but encountered the following error. Do you know what I should do?

@JeremyGe07 With onnx opset version 18, onnx.ReduceMean switched from using axes as an attribute to using axes as an optional input tensor. See onnx docs. You could try exporting with an opset version < 18 (as long as you don't need the support from newer opset versions).

Another thing that might be helpful if you want the most recent changes is to try using a nightly build of torch with

pip install torch -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html

I'd recommend submitting an issue for the ReduceMean problem, since this PR is likely not the best place to discuss this.

[justinchuby]

If you intend to export fx graphs to ONNX. I recommend using the torch.onnx.dynamo_export api. Bitwise ops should already be supported

[zjgarvey]

If you intend to export fx graphs to ONNX. I recommend using the torch.onnx.dynamo_export api. Bitwise ops should already be supported

Ah, that does seem to be the consensus. I'm submitting this PR since I'm working on projects which use torch.onnx.export. I'll look into simply changing over to torch.onnx.dynamo_export for those projects.

[justinchuby]

The pr itself looks good. Some tests may be preferable.

[justinchuby]

@pytorchbot rebase

[pytorchmergebot]

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

[pytorchmergebot]

Successfully rebased bitwise_ops_onnx_export onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via git checkout bitwise_ops_onnx_export && git pull --rebase)

[justinchuby]

@pytorchbot merge

[pytorchmergebot]

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

[zjgarvey]

The pr itself looks good. Some tests may be preferable.

Good point. I'm happy to add some tests (particularly for aten::bitwise_and) if you would like.

[pytorchmergebot]

Merge failed

Reason: 1 mandatory check(s) failed. The first few are:

• pull / linux-docs / build-docs-python-false

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

[titaiwangms]

Expose op to line 31

[titaiwangms]

Same for the other two ops

[zjgarvey]

I have a fix for most of the CI issues, but I'm not sure how to resolve the docs error:

undocumented objects found:
+ cat build/coverage/python.txt
Undocumented Python objects
===========================
torch.onnx.symbolic_opset18
---------------------------
Functions:
 * bitwise_and

+ echo 'Make sure you'\''ve updated relevant .rsts in docs/source!'
Make sure you've updated relevant .rsts in docs/source!
+ echo 'You can reproduce locally by running '\''cd docs && make coverage && cat build/coverage/python.txt'\'''
You can reproduce locally by running 'cd docs && make coverage && cat build/coverage/python.txt'

I'm able to reproduce the error, but I don't know what .rst to edit. Can someone give me advice on this?

[titaiwangms]

I'm able to reproduce the error, but I don't know what .rst to edit. Can someone give me advice on this?

I thnk it should be this part that we have to add the new op that we use:

pytorch/torch/onnx/symbolic_opset18.py

Lines 7 to 17 in 85fd76f

	New operators:
	CenterCropPad
	Col2Im
	Mish
	OptionalGetElement
	OptionalHasElement
	Pad
	Resize
	ScatterElements
	ScatterND
	Split

I reference on this: https://github.com/pytorch/pytorch/pull/118828/files#diff-4e77adb023a4e15f75009c08294251dfe9c3b7bb6947e7f7b4609c65edb53b20

[titaiwangms]

@pytorchbot merge

[pytorchmergebot]

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

[titaiwangms]

@zjgarvey Thank you!