Update compile and export tutorials for 2.1

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