pytorch

558e5a24 - Introduce torch.onnx.dynamo_export API (#97920)

Introduce torch.onnx.dynamo_export API (#97920) This is the first phase of the new ONNX exporter API for exporting from TorchDynamo and FX, and represents the beginning of a new era for exporting ONNX from PyTorch. The API here is a starting point upon which we will layer more capability and expressiveness in subsequent phases. This first phase introduces the following into `torch.onnx`: ```python dynamo_export( model: torch.nn.Module, /, *model_args, export_options: Optional[ExportOptions] = None, **model_kwargs, ) -> ExportOutput: ... class ExportOptions: opset_version: Optional[int] = None dynamic_shapes: Optional[bool] = None logger: Optional[logging.Logger] = None class ExportOutputSerializer(Protocol): def serialize( self, export_output: ExportOutput, destination: io.BufferedIOBase, ) -> None: ... class ExportOutput: model_proto: onnx.ModelProto def save( self, destination: Union[str, io.BufferedIOBase], *, serializer: Optional[ExportOutputSerializer] = None, ) -> None: ... ``` In addition to the API in the first commit on this PR, we have a few experiments for exporting Dynamo and FX to ONNX that this PR rationalizes through the new Exporter API and adjusts tests to use the new API. - A base `FXGraphModuleExporter` exporter from which all derive: - `DynamoExportExporter`: uses dynamo.export to acquire FX graph - `DynamoOptimizeExporter`: uses dynamo.optimize to acquire FX graph - `FXSymbolicTraceExporter`: uses FX symbolic tracing The `dynamo_export` API currently uses `DynamoOptimizeExporter`. ### Next Steps (subsequent PRs): * Combine `DynamoExportExporter` and `DynamoOptimizeExporter` into a single `DynamoExporter`. * Make it easy to test `FXSymbolicTraceExporter` through the same API; eventually `FXSymbolicTraceExporter` goes away entirely when the Dynamo approach works for large models. We want to keep `FXSymbolicTraceExporter` around for now for experimenting and internal use. * Parameterize (on `ExportOptions`) and consolidate Dynamo exporter tests. - This PR intentionally leaves the existing tests unchanged as much as possible except for the necessary plumbing. * Subsequent API phases: - Diagnostics - Registry, dispatcher, and Custom Ops - Passes - Dynamic shapes Fixes #94774 Pull Request resolved: https://github.com/pytorch/pytorch/pull/97920 Approved by: https://github.com/justinchuby, https://github.com/titaiwangms, https://github.com/thiagocrepaldi, https://github.com/shubhambhokare1