[ONNX] Single function for torch.onnx.export and torch.onnx.dynamo_export

[titaiwangms]

Add dynamo: bool = True as a switch in torch.onnx.export to provide users an option to try torch.onnx.dynamo_export.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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.

✅ You can merge normally! (1 Unrelated Failure)

As of commit af535f0 with merge base 22964d1 ():

UNSTABLE - The following job failed but was likely due to flakiness present on trunk and has been marked as unstable:

• pull / linux-focal-cuda12.4-py3.10-gcc9-sm86 / test (default, 2, 5, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) ()
  inductor/test_efficient_conv_bn_eval.py::EfficientConvBNEvalCudaTests::test_basic_cuda

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

[justinchuby]

LGTM thanks!

[gramalingam]

Can you say something about how this equality assertion works? Just for my own education, wondering what kind of infrastructure exists for proto-object equality check.

[titaiwangms]

That's a good point. I didn't verify that. This might be invalid. I will verify this.

[justinchuby]

There is some internal equality check according to https://proto-plus-python.readthedocs.io/en/stable/_modules/proto/message.html. I think it should work when two messages are exactly the same.

[titaiwangms]

It does catch the difference between two protobuf if there is any. Here is their code:

pytorch/torch/testing/_internal/common_utils.py

Line 3522 in 22964d1

def assertEqual(

[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

[justinchuby]

In a follow up PR, I would add an kwarg parameter to the api and encourage users to use that.

[titaiwangms]

torch.onnx.export has its own way to support kwargs. Do we want to add another way of supporting kwargs?

pytorch/torch/onnx/utils.py

Lines 240 to 281 in 4adee71

	3. A TUPLE OF ARGUMENTS ENDING WITH A DICTIONARY OF NAMED ARGUMENTS::
	args = (
	x,
	{
	"y": input_y,
	"z": input_z
	}
	)
	All but the last element of the tuple will be passed as non-keyword arguments,
	and named arguments will be set from the last element. If a named argument is
	not present in the dictionary, it is assigned the default value, or None if a
	default value is not provided.
	.. note::
	If a dictionary is the last element of the args tuple, it will be
	interpreted as containing named arguments. In order to pass a dict as the
	last non-keyword arg, provide an empty dict as the last element of the args
	tuple. For example, instead of::
	torch.onnx.export(
	model,
	(
	x,
	# WRONG: will be interpreted as named arguments
	{y: z}
	),
	"test.onnx.pb"
	)
	Write::
	torch.onnx.export(
	model,
	(
	x,
	{y: z},
	{}
	),
	"test.onnx.pb"
	)

[justinchuby]

Suggested change

	export_params: bool = True,
	*,
	export_params: bool = True,

as a follow up

[justinchuby]

Suggested change

	autograd_inlining: Optional[bool] = True,
	autograd_inlining: Optional[bool] = True,
	kwargs: Mapping[str, Any] | None = None,

As a follow up