[ONNX] Update API to torch.onnx.export(..., dynamo=True)

beginner_source/onnx/README.txt

@@ -10,5 +10,9 @@ ONNX
     https://pytorch.org/tutorials/beginner/onnx/export_simple_model_to_onnx_tutorial.html
 
 3. onnx_registry_tutorial.py
-    Extending the ONNX Registry
+    Extending the ONNX exporter operator support
     https://pytorch.org/tutorials/beginner/onnx/onnx_registry_tutorial.html
+
+4. export_control_flow_model_to_onnx_tutorial.py
+    Export a model with control flow to ONNX
+    https://pytorch.org/tutorials/beginner/onnx/export_control_flow_model_to_onnx_tutorial.html

beginner_source/onnx/export_control_flow_model_to_onnx_tutorial.py

@@ -0,0 +1,171 @@
+# -*- coding: utf-8 -*-
+"""
+`Introduction to ONNX <intro_onnx.html>`_ ||
+`Exporting a PyTorch model to ONNX <export_simple_model_to_onnx_tutorial.html>`_ ||
+`Extending the ONNX exporter operator support <onnx_registry_tutorial.html>`_ ||
+**`Export a model with control flow to ONNX**
+
+Export a model with control flow to ONNX
+==========================================
+
+**Author**: `Xavier Dupré <https://github.com/xadupre>`_.
+"""
+
+
+###############################################################################
+# Overview
+# --------
+# 
+# This tutorial demonstrates how to handle control flow logic while exporting
+# a PyTorch model to ONNX. It highlights the challenges of exporting
+# conditional statements directly and provides solutions to circumvent them.
+#
+# Conditional logic cannot be exported into ONNX unless they refactored
+# to use :func:`torch.cond`. Let's start with a simple model
+# implementing a test.
+
+import torch
+
+###############################################################################
+# Define the Models
+# --------
+#
+# Two models are defined:
+#
+# ForwardWithControlFlowTest: A model with a forward method containing an
+# if-else conditional.
+#
+# ModelWithControlFlowTest: A model that incorporates ForwardWithControlFlowTest
+# as part of a simple multi-layer perceptron (MLP). The models are tested with
+# a random input tensor to confirm they execute as expected.
+
+class ForwardWithControlFlowTest(torch.nn.Module):
+    def forward(self, x):
+        if x.sum():
+            return x * 2
+        return -x
+
+
+class ModelWithControlFlowTest(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.mlp = torch.nn.Sequential(
+            torch.nn.Linear(3, 2),
+            torch.nn.Linear(2, 1),
+            ForwardWithControlFlowTest(),
+        )
+
+    def forward(self, x):
+        out = self.mlp(x)
+        return out
+
+
+model = ModelWithControlFlowTest()
+
+
+###############################################################################
+# Exporting the Model: First Attempt
+# --------
+#
+# Exporting this model using torch.export.export fails because the control
+# flow logic in the forward pass creates a graph break that the exporter cannot
+# handle. This behavior is expected, as conditional logic not written using
+# torch.cond is unsupported.
+# 
+# A try-except block is used to capture the expected failure during the export
+# process. If the export unexpectedly succeeds, an AssertionError is raised.
+
+x = torch.randn(3)
+model(x)
+
+try:
+    torch.export.export(model, (x,), strict=False)
+    raise AssertionError("This export should failed unless PyTorch now supports this model.")
+except Exception as e:
+    print(e)
+
+###############################################################################
+# Using torch.onnx.export with JIT Tracing
+# --------
+#
+# When exporting the model using torch.onnx.export with the dynamo=True
+# argument, the exporter defaults to using JIT tracing. This fallback allows
+# the model to export, but the resulting ONNX graph may not faithfully represent
+# the original model logic due to the limitations of tracing.
+
+
+onnx_program = torch.onnx.export(model, (x,), dynamo=True) 
+print(onnx_program.model)
+
+
+###############################################################################
+# Suggested Patch: Refactoring with torch.cond
+# --------
+#
+# To make the control flow exportable, the tutorial demonstrates replacing the
+# forward method in ForwardWithControlFlowTest with a refactored version that
+# uses torch.cond.
+#
+# Details of the Refactoring:
+#
+# Two helper functions (identity2 and neg) represent the branches of the conditional logic:
+# * torch.cond is used to specify the condition and the two branches along with the input arguments.
+# * The updated forward method is then dynamically assigned to the ForwardWithControlFlowTest instance within the model. A list of submodules is printed to confirm the replacement.
+
+def new_forward(x):
+    def identity2(x):
+        return x * 2
+
+    def neg(x):
+        return -x
+
+    return torch.cond(x.sum() > 0, identity2, neg, (x,))
+
+
+print("the list of submodules")
+for name, mod in model.named_modules():
+    print(name, type(mod))
+    if isinstance(mod, ForwardWithControlFlowTest):
+        mod.forward = new_forward
+
+###############################################################################
+# Let's see what the fx graph looks like.
+
+print(torch.export.export(model, (x,), strict=False))  
+
+###############################################################################
+# Let's export again.
+
+onnx_program = torch.onnx.export(model, (x,), dynamo=True)  
+print(onnx_program.model) 
+
+
+###############################################################################
+# We can optimize the model and get rid of the model local functions created to capture the control flow branches.  
+
+onnx_program.optimize()  
+print(onnx_program.model)  
+
+###############################################################################
+# Conclusion
+# --------
+# This tutorial demonstrates the challenges of exporting models with conditional
+# logic to ONNX and presents a practical solution using torch.cond.
+# While the default exporters may fail or produce imperfect graphs, refactoring the
+# model's logic ensures compatibility and generates a faithful ONNX representation.
+#
+# By understanding these techniques, we can overcome common pitfalls when
+# working with control flow in PyTorch models and ensure smooth integration with ONNX workflows.
+#
+# Further reading
+# ---------------
+#
+# The list below refers to tutorials that ranges from basic examples to advanced scenarios,
+# not necessarily in the order they are listed.
+# Feel free to jump directly to specific topics of your interest or
+# sit tight and have fun going through all of them to learn all there is about the ONNX exporter.
+#
+# .. include:: /beginner_source/onnx/onnx_toc.txt
+#
+# .. toctree::
+#    :hidden:

beginner_source/onnx/export_simple_model_to_onnx_tutorial.py

@@ -2,17 +2,18 @@
 """
 `Introduction to ONNX <intro_onnx.html>`_ ||
 **Exporting a PyTorch model to ONNX** ||
-`Extending the ONNX Registry <onnx_registry_tutorial.html>`_
+`Extending the ONNX exporter operator support <onnx_registry_tutorial.html>`_ ||
+`Export a model with control flow to ONNX <export_control_flow_model_to_onnx_tutorial.html>`_
 
 Export a PyTorch model to ONNX
 ==============================
 
-**Author**: `Ti-Tai Wang <https://github.com/titaiwangms>`_ and `Xavier Dupré <https://github.com/xadupre>`_
+**Author**: `Ti-Tai Wang <https://github.com/titaiwangms>`_ and Thiago Crepaldi <https://github.com/thiagocrepaldi>`_.
 
 .. note::
-    As of PyTorch 2.1, there are two versions of ONNX Exporter.
+    As of PyTorch 2.5, there are two versions of ONNX Exporter.
 
-    * ``torch.onnx.dynamo_export`` is the newest (still in beta) exporter based on the TorchDynamo technology released with PyTorch 2.0
+    * ``torch.onnx.export(..., dynamo=True)`` is the newest (still in beta) exporter based on the TorchDynamo technology released with PyTorch 2.0
     * ``torch.onnx.export`` is based on TorchScript backend and has been available since PyTorch 1.2.0
 
 """
@@ -21,7 +22,7 @@
 # In the `60 Minute Blitz <https://pytorch.org/tutorials/beginner/deep_learning_60min_blitz.html>`_,
 # we had the opportunity to learn about PyTorch at a high level and train a small neural network to classify images.
 # In this tutorial, we are going to expand this to describe how to convert a model defined in PyTorch into the
-# ONNX format using TorchDynamo and the ``torch.onnx.dynamo_export`` ONNX exporter.
+# ONNX format using TorchDynamo and the ``torch.onnx.export(..., dynamo=True)`` ONNX exporter.
 #
 # While PyTorch is great for iterating on the development of models, the model can be deployed to production
 # using different formats, including `ONNX <https://onnx.ai/>`_ (Open Neural Network Exchange)!
@@ -90,7 +91,16 @@ def forward(self, x):
 
 torch_model = MyModel()
 torch_input = torch.randn(1, 1, 32, 32)
-onnx_program = torch.onnx.dynamo_export(torch_model, torch_input)
+onnx_program = torch.onnx.export(torch_model, torch_input, dynamo=True)
+
+######################################################################
+# 3.5. (Optional) Optimize the ONNX model
+# ---------------------------------------
+#
+# The ONNX model can be optimized with constant folding, and elimination of redundant nodes.
+# The optimization is done in-place, so the original ONNX model is modified.
+
+onnx_program.optimize()
 
 ######################################################################
 # As we can see, we didn't need any code change to the model.

beginner_source/onnx/intro_onnx.py

@@ -1,13 +1,14 @@
 """
 **Introduction to ONNX** ||
 `Exporting a PyTorch model to ONNX <export_simple_model_to_onnx_tutorial.html>`_ ||
-`Extending the ONNX Registry <onnx_registry_tutorial.html>`_
+`Extending the ONNX exporter operator support <onnx_registry_tutorial.html>`_ ||
+`Export a model with control flow to ONNX <export_control_flow_model_to_onnx_tutorial.html>`_
 
 Introduction to ONNX
 ====================
 
 Authors:
-`Ti-Tai Wang <https://github.com/titaiwangms>`_ and `Xavier Dupré <https://github.com/xadupre>`_
+`Ti-Tai Wang <https://github.com/titaiwangms>`_ and Thiago Crepaldi <https://github.com/thiagocrepaldi>`_.
 
 `Open Neural Network eXchange (ONNX) <https://onnx.ai/>`_ is an open standard
 format for representing machine learning models. The ``torch.onnx`` module provides APIs to
@@ -19,8 +20,10 @@
 including Microsoft's `ONNX Runtime <https://www.onnxruntime.ai>`_.
 
 .. note::
-    Currently, there are two flavors of ONNX exporter APIs,
-    but this tutorial will focus on the ``torch.onnx.dynamo_export``.
+    Currently, the users can choose either through `TorchScript https://pytorch.org/docs/stable/jit.html`_ or
+    `ExportedProgram https://pytorch.org/docs/stable/export.html`_ to export the model to ONNX by the 
+    boolean parameter dynamo in `torch.onnx.export <https://pytorch.org/docs/stable/generated/torch.onnx.export.html>`_.
+    In this tutorial, we will focus on the ExportedProgram approach.
 
 The TorchDynamo engine is leveraged to hook into Python's frame evaluation API and dynamically rewrite its
 bytecode into an `FX graph <https://pytorch.org/docs/stable/fx.html>`_.
@@ -33,7 +36,7 @@
 Dependencies
 ------------
 
-PyTorch 2.1.0 or newer is required.
+PyTorch 2.5.0 or newer is required.
 
 The ONNX exporter depends on extra Python packages: