Enable strings in guess_dynamic_shapes

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>`_

_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?
 # ++++++++++++++++++++

_unittests/ut_export/test_dynamic_shapes.py

@@ -470,6 +470,65 @@ def forward(self, cache, z):
             ),
         )
 
+    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))

onnx_diagnostic/export/dynamic_shapes.py

@@ -723,13 +723,14 @@ def module_name_type(self):
         return f"type({self.name})={self.true_model_name}.{self.method_name}"
 
     def guess_dynamic_dimensions(
-        self, *tensors, auto: bool = False
+        self, *tensors, auto: Union[bool, str] = False
     ) -> Optional[Dict[int, Any]]:
         """
         Infers the dynamic dimension from multiple shapes.
         If auto is True, it returns ``torch.export.Dim.AUTO`` for every dimension
         which cannot be guessed. Two tensors with the same value for one dimension
-        can be guessed, but if there is only 1, it cannot.
+        can be guessed, but if there is only 1, it cannot. ``auto``` can be a string
+        to produce strings.
         """
         if len(tensors) == 1:
             if isinstance(tensors[0], (int, float)):
@@ -740,7 +741,7 @@ def guess_dynamic_dimensions(
             )
             return (
                 {i: torch.export.Dim.AUTO for i in range(len(tensors[0].shape))}  # noqa: C420
-                if auto
+                if auto and not isinstance(auto, str)
                 else {}
             )
         shapes = [t.shape for t in tensors]
@@ -750,22 +751,26 @@ def guess_dynamic_dimensions(
             f"shapes={shapes} for module {self.name!r}, "
             f"class={self.true_model_name!r}"
         )
-        dynamic: Any = torch.export.Dim.DYNAMIC  # type: ignore
+        dynamic: Any = (
+            auto
+            if isinstance(auto, str)
+            else (torch.export.Dim.AUTO if auto else torch.export.Dim.DYNAMIC)
+        )
         rk = set_length.pop()
         res = {}
         for i in range(rk):
             set_dim = set(s[i] for s in shapes)
             if len(set_dim) > 1:
-                res[i] = dynamic
+                res[i] = dynamic if not isinstance(dynamic, str) else f"{dynamic}{i}"
                 continue
             if set_dim == {0}:
                 # It is unexpected to find a null dimension. Let's replace it by a dynamic one.
-                res[i] = dynamic
+                res[i] = dynamic if not isinstance(dynamic, str) else f"{dynamic}{i}"
                 continue
         return res
 
     def guess_dynamic_shape_object(
-        self, *objs: Any, auto: bool = False, msg: Optional[Callable] = None
+        self, *objs: Any, auto: Union[bool, str] = False, msg: Optional[Callable] = None
     ) -> Any:
         """Guesses the dynamic shapes for one argument."""
         if len(objs) == 0:
@@ -790,7 +795,11 @@ def guess_dynamic_shape_object(
             shapes: Any = []
             for i in range(kl.pop()):
                 shapes.append(
-                    self.guess_dynamic_shape_object(*[o[i] for o in objs], auto=auto, msg=msg)
+                    self.guess_dynamic_shape_object(
+                        *[o[i] for o in objs],
+                        auto=auto if isinstance(auto, bool) else f"{auto}_{i}t",
+                        msg=msg,
+                    )
                 )
             return tuple(shapes)
 
@@ -802,7 +811,11 @@ def guess_dynamic_shape_object(
             shapes = []
             for i in range(kl.pop()):
                 shapes.append(
-                    self.guess_dynamic_shape_object(*[o[i] for o in objs], auto=auto, msg=msg)
+                    self.guess_dynamic_shape_object(
+                        *[o[i] for o in objs],
+                        auto=auto if isinstance(auto, bool) else f"{auto}_{i}l",
+                        msg=msg,
+                    )
                 )
             return shapes
 
@@ -814,7 +827,9 @@ def guess_dynamic_shape_object(
             shapes = {}
             for i in obj:
                 shapes[i] = self.guess_dynamic_shape_object(
-                    *[o[i] for o in objs], auto=auto, msg=msg
+                    *[o[i] for o in objs],
+                    auto=auto if isinstance(auto, bool) else f"{auto}_{i}d",
+                    msg=msg,
                 )
             return shapes
 
@@ -834,7 +849,9 @@ def guess_dynamic_shape_object(
             for i in range(kc.pop()):
                 values.append(
                     self.guess_dynamic_shape_object(
-                        *[ca[i] for ca in col_args], auto=auto, msg=msg
+                        *[ca[i] for ca in col_args],
+                        auto=auto if isinstance(auto, bool) else f"{auto}_{i}o",
+                        msg=msg,
                     )
                 )
             return values
@@ -852,12 +869,18 @@ def guess_dynamic_shape_object(
             key_cache = []
             for i in range(kc.pop()):
                 key_cache.append(
-                    self.guess_dynamic_dimensions(*[o.key_cache[i] for o in objs], auto=auto)
+                    self.guess_dynamic_dimensions(
+                        *[o.key_cache[i] for o in objs],
+                        auto=auto if isinstance(auto, bool) else f"{auto}_{i}kdc",
+                    )
                 )
             value_cache = []
             for i in range(vc.pop()):
                 value_cache.append(
-                    self.guess_dynamic_dimensions(*[o.value_cache[i] for o in objs], auto=auto)
+                    self.guess_dynamic_dimensions(
+                        *[o.value_cache[i] for o in objs],
+                        auto=auto if isinstance(auto, bool) else f"{auto}_{i}vdc",
+                    )
                 )
             return [key_cache, value_cache]
 
@@ -867,13 +890,17 @@ def guess_dynamic_shape_object(
             f"this object needs serialization function to be registered."
         )
 
-    def guess_dynamic_shapes(self, auto: bool = False) -> DYNAMIC_SHAPES:
+    def guess_dynamic_shapes(self, auto: Union[bool, str] = False) -> DYNAMIC_SHAPES:
         """
         Guesses the dynamic shapes for that module from two execution.
         If there is only one execution, then that would be static dimensions.
 
         :param auto: if auto is True, use ``torch.export.Dim.AUTO`` for any
-            dimension if the number of inputs is one
+            dimension if the number of inputs is one,
+            if ``auto`` is a string, it uses strings
+        :return: guessed dynamic shapes
+
+        See example :ref:`l-guess-dynamic-shapes-example`.
         """
         if len(self.inputs) == 0:
             # No inputs, unable to guess.
@@ -900,7 +927,9 @@ def guess_dynamic_shapes(self, auto: bool = False) -> DYNAMIC_SHAPES:
             objs = [_[0][i] for _ in self.inputs]
             args.append(
                 self.guess_dynamic_shape_object(
-                    *objs, auto=auto, msg=lambda i=i: f" failing input {i}"
+                    *objs,
+                    auto=auto if isinstance(auto, bool) else f"{auto}_{i}I",
+                    msg=lambda i=i: f" failing input {i}",
                 )
             )
         names = s2.pop()
@@ -913,7 +942,9 @@ def guess_dynamic_shapes(self, auto: bool = False) -> DYNAMIC_SHAPES:
 
             objs = [_[1][name] for _ in self.inputs]
             kwargs[name] = self.guess_dynamic_shape_object(
-                *objs, auto=auto, msg=lambda name=name: f" failing input {name!r}"
+                *objs,
+                auto=auto if isinstance(auto, bool) else f"{auto}_{i}I",
+                msg=lambda name=name: f" failing input {name!r}",
             )
         return tuple(args), kwargs