fix doc

Author: xadupre

_doc/conf.py

@@ -233,6 +233,7 @@ def linkcode_resolve(domain, info):
     "onnx-script": "https://github.com/microsoft/onnxscript",
     "onnxscript": "https://github.com/microsoft/onnxscript",
     "onnxscript Tutorial": "https://microsoft.github.io/onnxscript/tutorial/index.html",
+    "optree": "https://github.com/metaopt/optree",
     "Pattern-based Rewrite Using Rules With onnxscript": "https://microsoft.github.io/onnxscript/tutorial/rewriter/rewrite_patterns.html",
     "opsets": "https://onnx.ai/onnx/intro/concepts.html#what-is-an-opset-version",
     "pyinstrument": "https://pyinstrument.readthedocs.io/en/latest/",

_doc/recipes/plot_dynamic_shapes_json.py

@@ -2,8 +2,17 @@
 JSON returns list when the original dynamic shapes are list or tuple
 ====================================================================
 
+Dynamic shapes given to :func:`torch.export.export` must follow the
+same semantic. What if we confuse tuple and list when defining the dynamic shapes,
+how to restore the expected type assuming we know the inputs?
+Not often useful but maybe we will learn more about
+:epkg:`optree`.
+
 Dynamic Shapes After JSON
 +++++++++++++++++++++++++
+
+JSON format does not make the difference between a list and a tuple.
+So after serializing to json and restoring, both of them become lists.
 """
 
 import json
@@ -49,8 +58,8 @@
 # %%
 # tuple are replaced by list.
 
-# The trick
-# +++++++++
+# The trick to restore tuple when expected
+# ++++++++++++++++++++++++++++++++++++++++
 
 
 def flatten_unflatten_like_dynamic_shapes(obj):

_unittests/ut_export/test_shape_helper.py

@@ -10,23 +10,24 @@ class TestShapeHelper(ExtTestCase):
     @requires_torch("2.7.99")
     def test_all_dynamic_shape_from_inputs(self):
         ds = all_dynamic_shape_from_inputs((torch.randn((5, 6)), torch.randn((1, 6))))
+        self.assertEqual(({0: "d_0_0", 1: "d_0_1"}, {0: "d_1_0", 1: "d_1_1"}), ds)
+        ds = all_dynamic_shape_from_inputs([torch.randn((5, 6)), torch.randn((1, 6))])
         self.assertEqual([{0: "d_0_0", 1: "d_0_1"}, {0: "d_1_0", 1: "d_1_1"}], ds)
         ds = all_dynamic_shape_from_inputs(
             (torch.randn((5, 6)), torch.randn((1, 6))), dim_prefix=torch.export.Dim.AUTO
         )
         self.assertEqual(
-            [
+            (
                 {0: torch.export.Dim.AUTO, 1: torch.export.Dim.AUTO},
                 {0: torch.export.Dim.AUTO, 1: torch.export.Dim.AUTO},
-            ],
+            ),
             ds,
         )
 
     @requires_transformers("4.52")
     @requires_torch("2.7.99")
     def test_all_dynamic_shape_from_inputs_dynamic_cache(self):
         data = get_untrained_model_with_inputs("arnir0/Tiny-LLM")
-        print(self.string_type(data["inputs"], with_shape=True))
         ds = all_dynamic_shape_from_inputs(data["inputs"])
         self.assertEqual(
             {