fix doc

Author: xadupre

_doc/conf.py

@@ -126,6 +126,7 @@
     ("py:class", "transformers.configuration_utils.PretrainedConfig"),
     ("py:func", "torch.export._draft_export.draft_export"),
     ("py:func", "torch._export.tools.report_exportability"),
+    ("py:func", "torch.utils._pytree.register_pytree_node"),
     ("py:meth", "huggingface_hub.HfApi.list_models"),
     ("py:meth", "transformers.AutoConfig.from_pretrained"),
     ("py:meth", "transformers.GenerationMixin.generate"),

_doc/examples/plot_export_with_args_kwargs.py

@@ -1,20 +1,16 @@
 """
 .. _l-plot-export-with-args-kwargs:
 
-==================================
-Dynamic Shapes for *args, **kwargs
-==================================
+==========================================
+Dynamic Shapes for ``*args``, ``**kwargs``
+==========================================
 
 Quick tour of dynamic shapes.
-
-Simple Examples
-===============
-
 We first look at examples playing positional and names parameters
 to understand how :func:`torch.export.export` works.
 
 args
-++++
+====
 """
 
 import pprint
@@ -59,7 +55,7 @@ def forward(self, x, y):
 
 # %%
 # kwargs
-# ++++++
+# ======
 #
 # We do the same with named arguments.
 
@@ -88,7 +84,7 @@ def forward(self, x, y):
 
 # %%
 # args and kwargs
-# +++++++++++++++
+# ===============
 #
 # :func:`torch.export.export` does not like having dynami shapes
 # for both args and kwargs. We need to define them using one mechanism.