Merge branch 'main' into fix_render

Author: Svetlana Karslioglu

.jenkins/build.sh

@@ -25,9 +25,9 @@ pip install --progress-bar off -r $DIR/../requirements.txt
 #Install PyTorch Nightly for test.
 # Nightly - pip install --pre torch torchvision torchaudio -f https://download.pytorch.org/whl/nightly/cu102/torch_nightly.html
 # Install 2.1 for testing
-pip uninstall -y torch torchvision torchaudio torchtext torchdata
-pip3 install torch torchvision torchaudio --no-cache-dir --index-url https://download.pytorch.org/whl/test/cu121
-pip3 install torchdata torchtext --index-url https://download.pytorch.org/whl/test/cpu
+# pip uninstall -y torch torchvision torchaudio torchtext torchdata
+# pip3 install torch torchvision torchaudio --no-cache-dir --index-url https://download.pytorch.org/whl/test/cu121
+# pip3 install torchdata torchtext --index-url https://download.pytorch.org/whl/test/cpu
 
 # Install two language tokenizers for Translation with TorchText tutorial
 python -m spacy download en_core_web_sm

intermediate_source/torch_compile_tutorial.py

@@ -20,11 +20,8 @@
 #
 # **Contents**
 #
-# - Basic Usage
-# - Demonstrating Speedups
-# - Comparison to TorchScript and FX Tracing
-# - TorchDynamo and FX Graphs
-# - Conclusion
+# .. contents::
+#     :local:
 #
 # **Required pip Dependencies**
 #
@@ -485,19 +482,12 @@ def bar(a, b):
 print(opt_model(generate_data(16)[0]))
 
 ######################################################################
-# <!----TODO: replace this section with a link to the torch.export tutorial when done --->
-#
-# Finally, if we simply want TorchDynamo to output the FX graph for export,
-# we can use ``torch._dynamo.export``. Note that ``torch._dynamo.export``, like
-# ``fullgraph=True``, raises an error if TorchDynamo breaks the graph.
-
-try:
-    torch._dynamo.export(bar)(torch.randn(10), torch.randn(10))
-except:
-    tb.print_exc()
-
-model_exp = torch._dynamo.export(init_model())(generate_data(16)[0])
-print(model_exp[0](generate_data(16)[0]))
+# We can use ``torch.export`` (from PyTorch 2.1+) to extract a single, exportable
+# FX graph from the input PyTorch program. The exported graph is intended to be
+# run on different (i.e. Python-less) environments. One important restriction
+# is that the ``torch.export`` does not support graph breaks. Please check
+# `this tutorial <https://pytorch.org/tutorials/intermediate/torch_export_tutorial.html>`__
+# for more details on ``torch.export``.
 
 ######################################################################
 # Conclusion

intermediate_source/torch_export_tutorial.py

@@ -423,7 +423,7 @@ def constraints_example4(x, y):
 exported_constraints_example4 = export(constraints_example4, (torch.randn(3, 3), torch.tensor([4])))
 print(exported_constraints_example4(torch.randn(3, 3), torch.tensor([5])))
 try:
-    exported_constraints_example4(torch.randn(3, 3), torch.randn([2]))
+    exported_constraints_example4(torch.randn(3, 3), torch.tensor([2]))
 except Exception:
     tb.print_exc()
 
@@ -441,7 +441,7 @@ def constraints_example5(x, y):
 exported_constraints_example5 = export(constraints_example5, (torch.randn(2, 2), torch.tensor([4])))
 print(exported_constraints_example5(torch.randn(2, 2), torch.tensor([5])))
 try:
-    exported_constraints_example5(torch.randn(2, 2), torch.randn([1]))
+    exported_constraints_example5(torch.randn(2, 2), torch.tensor([1]))
 except Exception:
     tb.print_exc()
 
@@ -496,6 +496,10 @@ def custom_op_example(x):
 # Note in the above outputs that the custom op is included in the exported graph.
 # And when we call the exported graph as a function, the original custom op is called,
 # as evidenced by the ``print`` call.
+#
+# If you have a custom operator implemented in C++, please refer to
+# `this document <https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit#heading=h.ahugy69p2jmz>`__
+# to make it compatible with ``torch.export``.
 
 ######################################################################
 # ExportDB