Enable strings in guess_dynamic_shapes (#91)

* Enable strings in guess_dynamic_shapes

* fix

* doc

Author: xadupre

README.rst

@@ -67,7 +67,7 @@ Enlightening Examples
   <https://sdpython.github.io/doc/onnx-diagnostic/dev/auto_examples/plot_export_with_dynamic_shapes_auto.html>`_
 * `Find and fix an export issue due to dynamic shapes
   <https://sdpython.github.io/doc/onnx-diagnostic/dev/auto_examples/plot_export_locate_issue.html>`_
-* `Export with DynamicCache and dynamic shapes
+* `Export with DynamicCache and guessed dynamic shapes
   <https://sdpython.github.io/doc/onnx-diagnostic/dev/auto_examples/plot_export_with_dynamic_cache.html>`_
 * `Steel method forward to guess the dynamic shapes (with Tiny-LLM)
   <https://sdpython.github.io/doc/onnx-diagnostic/dev/auto_examples/plot_export_tiny_llm.html>`_
@@ -95,10 +95,7 @@ Snapshot of usefuls tools
 
     inputs = (
         torch.rand((3, 4), dtype=torch.float16),
-        [
-            torch.rand((5, 6), dtype=torch.float16),
-            torch.rand((5, 6, 7), dtype=torch.float16),
-        ]
+        [torch.rand((5, 6), dtype=torch.float16), torch.rand((5, 6, 7), dtype=torch.float16)],
     )
 
     # with shapes
@@ -126,4 +123,33 @@ Snapshot of usefuls tools
 
 **max_diff**
 
-Returns the maximum discrancies across nested containers containing tensors.
+.. code-block:: python
+
+    import torch
+    from onnx_diagnostic.helpers import max_diff
+
+    print(
+        max_diff(
+            (torch.Tensor([1, 2]), (torch.Tensor([1, 2]),)),
+            (torch.Tensor([1, 2]), (torch.Tensor([1, 2]),)),
+        )
+    )
+
+::
+
+    >>> {"abs": 0.0, "rel": 0.0, "sum": 0.0, "n": 4.0, "dnan": 0.0}s
+
+**guess_dynamic_shapes**
+
+.. code-block:: python
+
+    inputs = [
+        (torch.randn((5, 6)), torch.randn((1, 6))),
+        (torch.randn((7, 8)), torch.randn((1, 8))),
+    ]
+    ds = ModelInputs(model, inputs).guess_dynamic_shapes(auto="dim")
+    print(ds)
+
+::
+
+    >>> (({0: 'dim_0I0', 1: 'dim_0I1'}, {1: 'dim_1I1'}), {})

_doc/examples/plot_export_with_dynamic_cache.py

@@ -1,9 +1,9 @@
 """
 .. _l-plot-export-with-dynamic-shape:
 
-===========================================
-Export with DynamicCache and dynamic shapes
-===========================================
+===================================================
+Export with DynamicCache and guessed dynamic shapes
+===================================================
 
 Every LLMs implemented in :epkg:`transformers` use cache.
 One of the most used is :class:`transformers.cache_utils.DynamicCache`.
@@ -84,6 +84,8 @@ def forward(self, cache, z):
 print(string_type(inputs[1], with_shape=True))
 
 # %%
+# .. _l-guess-dynamic-shapes-example:
+#
 # Guess the dynamic shapes
 # ========================
 #
@@ -112,6 +114,17 @@ def forward(self, cache, z):
         )
 print(ep)
 
+# %%
+# Use string instead of DYNAMIC
+# +++++++++++++++++++++++++++++
+#
+# ONNX exporter considers strings instead of DYNAMIC or AUTO
+# to give names to every dimension.
+
+dss = mi.guess_dynamic_shapes(auto="dim")
+pprint.pprint(dss)
+
+
 # %%
 # Do we need to guess?
 # ++++++++++++++++++++

_doc/index.rst

@@ -1,4 +1,4 @@
-
+========================================
 onnx-diagnostic: investigate onnx models
 ========================================
 
@@ -48,7 +48,7 @@ It also implements tools to investigate, validate exported models (ExportedProgr
     license
 
 Getting started
-+++++++++++++++
+===============
 
 ::
 
@@ -63,7 +63,7 @@ or
     pip install onnx-diagnostic
 
 Enlightening Examples
-+++++++++++++++++++++
+=====================
 
 **Where to start to export a model**
 
@@ -85,7 +85,13 @@ Enlightening Examples
 * :ref:`l-plot-failing-onnxruntime-evaluator`
 * :ref:`l-plot-failing-model-extract`
 
-**Some Usefuls Tools**
+Some Usefuls Tools
+==================
+
+string_type
++++++++++++
+
+See :func:`onnx_diagnostic.helpers.string_type`.
 
 .. code-block:: python
 
@@ -107,6 +113,11 @@ Enlightening Examples
 
     >>> (T10s3x4,#2[T10s5x6,T10s5x6x7])
 
+onnx_dtype_name
++++++++++++++++
+
+See :func:`onnx_diagnostic.helpers.onnx_dtype_name`.
+
 .. code-block:: python
 
         import onnx
@@ -121,7 +132,64 @@ Enlightening Examples
     >>> BFLOAT16
     >>> INT64
 
-:func:`onnx_diagnostic.helpers.max_diff`, ...
+max_diff
+++++++++
+
+See :func:`onnx_diagnostic.helpers.max_diff`.
+
+.. code-block:: python
+
+    import torch
+    from onnx_diagnostic.helpers import max_diff
+
+    print(
+        max_diff(
+            (torch.Tensor([1, 2]), (torch.Tensor([1, 2]),)),
+            (torch.Tensor([1, 2]), (torch.Tensor([1, 2]),)),
+        )
+    )
+
+::
+
+    >>> {"abs": 0.0, "rel": 0.0, "sum": 0.0, "n": 4.0, "dnan": 0.0}s
+
+guess_dynamic_shapes
+++++++++++++++++++++
+
+See :meth:`onnx_diagnostic.export.ModelInputs.guess_dynamic_shapes`.
+
+.. code-block:: python
+
+    inputs = [
+        (torch.randn((5, 6)), torch.randn((1, 6))),
+        (torch.randn((7, 8)), torch.randn((1, 8))),
+    ]
+    ds = ModelInputs(model, inputs).guess_dynamic_shapes(auto="dim")
+    print(ds)
+
+::
+
+    >>> (({0: 'dim_0I0', 1: 'dim_0I1'}, {1: 'dim_1I1'}), {})
+
+use_dyn_for_str
++++++++++++++++
+
+
+
+Older versions
+++++++++++++++
+
+* `0.5.0 <../v0.5.0/index.html>`_
+* `0.4.4 <../v0.4.4/index.html>`_
+* `0.4.3 <../v0.4.3/index.html>`_
+* `0.4.2 <../v0.4.2/index.html>`_
+* `0.4.1 <../v0.4.1/index.html>`_
+* `0.4.0 <../v0.4.0/index.html>`_
+* `0.3.0 <../v0.3.0/index.html>`_
+* `0.2.2 <../v0.2.2/index.html>`_
+* `0.2.1 <../v0.2.1/index.html>`_
+* `0.2.0 <../v0.2.0/index.html>`_
+* `0.1.0 <../v0.1.0/index.html>`_
 
 The documentation was updated on:
 
@@ -173,18 +241,3 @@ Size of the package:
     df = pandas.DataFrame(statistics_on_folder(os.path.dirname(__file__), aggregation=1))
     gr = df[["dir", "ext", "lines", "chars"]].groupby(["ext", "dir"]).sum()
     print(gr)
-
-Older versions
-++++++++++++++
-
-* `0.5.0 <../v0.5.0/index.html>`_
-* `0.4.4 <../v0.4.4/index.html>`_
-* `0.4.3 <../v0.4.3/index.html>`_
-* `0.4.2 <../v0.4.2/index.html>`_
-* `0.4.1 <../v0.4.1/index.html>`_
-* `0.4.0 <../v0.4.0/index.html>`_
-* `0.3.0 <../v0.3.0/index.html>`_
-* `0.2.2 <../v0.2.2/index.html>`_
-* `0.2.1 <../v0.2.1/index.html>`_
-* `0.2.0 <../v0.2.0/index.html>`_
-* `0.1.0 <../v0.1.0/index.html>`_

_unittests/ut_export/test_dynamic_shapes.py

@@ -470,6 +470,66 @@ def forward(self, cache, z):
             ),
         )
 
+    @requires_transformers("4.51")
+    def test_guess_dynamic_shapes_cache_str(self):
+        class Model(torch.nn.Module):
+            def forward(self, cache, z):
+                return (
+                    z
+                    + cache.key_cache[0]
+                    + cache.key_cache[1]
+                    + cache.value_cache[0]
+                    + cache.value_cache[1]
+                )
+
+        model = Model()
+
+        n_layers = 2
+        bsize, nheads, slen, dim = 2, 4, 3, 7
+        cache = make_dynamic_cache(
+            [
+                (torch.randn(bsize, nheads, slen, dim), torch.randn(bsize, nheads, slen, dim))
+                for i in range(n_layers)
+            ]
+        )
+        z = torch.randn((1, 1, 1, 7))
+        model(cache, z)
+
+        cache2 = make_dynamic_cache(
+            [
+                (
+                    torch.randn(bsize, nheads, slen, dim),
+                    torch.randn(bsize + 1, nheads, slen + 1, dim + 1),
+                )
+                for i in range(n_layers)
+            ]
+        )
+        inputs = [
+            (cache, z),
+            (cache2, torch.randn((1, 1, 1, 8))),
+        ]
+
+        mi = ModelInputs(Model(), inputs)
+        self.assertIn("DynamicCache", string_type(mi.inputs, with_shape=True))
+        ds = mi.guess_dynamic_shapes(auto="dim")
+        print(ds)
+        self.assertEqual(
+            ds,
+            (
+                (
+                    [
+                        [{}, {}],
+                        [
+                            {0: "dim_0I_1o_0l0", 2: "dim_0I_1o_0l2", 3: "dim_0I_1o_0l3"},
+                            {0: "dim_0I_1o_1l0", 2: "dim_0I_1o_1l2", 3: "dim_0I_1o_1l3"},
+                        ],
+                    ],
+                    {3: "dim_1I3"},
+                ),
+                {},
+            ),
+        )
+
     def test_couple_input_ds_0(self):
         T3x4 = torch.rand((3, 4))
         T3x1 = torch.rand((3, 1))