changes

Author: xadupre

CHANGELOGS.rst

@@ -4,6 +4,10 @@ Change Logs
 0.4.0
 +++++
 
+* :pr:`58`: add function use_dyn_not_str to replace string by ``torch.export.Dim.DYNAMIC``,
+  use string instead of ``torch.export.Dim.DYNAMIC`` when returning the dynamic shapes
+  for a specific models, it is a valid definition for ``torch.onnx.export``
+  which can reuse the names
 * :pr:`55`: add support for text-classification
 * :pr:`54`: add support for fill-mask, refactoring
 * :pr:`52`: add support for zero-shot-image-classification

_doc/examples/plot_export_hub_codellama.py

@@ -30,6 +30,7 @@
     task_from_id,
 )
 from onnx_diagnostic.torch_export_patches import bypass_export_some_errors
+from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
 
 model_id = "codellama/CodeLlama-7b-Python-hf"
 print("info", get_model_info(model_id))
@@ -96,7 +97,7 @@
         model,
         (),
         kwargs=f(data["inputs"]),
-        dynamic_shapes=data["dynamic_shapes"],
+        dynamic_shapes=use_dyn_not_str(data["dynamic_shapes"]),
         strict=False,
     )
     print(ep)

_doc/examples/plot_export_tiny_phi2.py

@@ -27,6 +27,7 @@
 from onnx_diagnostic.helpers.cache_helper import is_cache_dynamic_registered
 from onnx_diagnostic.helpers.rt_helper import make_feeds
 from onnx_diagnostic.torch_export_patches import bypass_export_some_errors
+from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
 from onnx_diagnostic.torch_models.hghub import (
     get_untrained_model_with_inputs,
 )
@@ -92,7 +93,7 @@
         untrained_model,
         (),
         kwargs=modificator(copy.deepcopy(inputs)),
-        dynamic_shapes=dynamic_shapes,
+        dynamic_shapes=use_dyn_not_str(dynamic_shapes),
         strict=False,  # mandatory for torch==2.6
     )
 

_unittests/ut_export/test_dynamic_shapes.py

@@ -576,17 +576,20 @@ def test_couple_input_ds_cache(self):
                 Cls(
                     (),
                     kwargs,
-                    {"A": ds_batch, "B": (ds_batch, [ds_batch, ds_batch, ds_batch, ds_batch])},
+                    {
+                        "A": ds_batch,
+                        "B": (ds_batch, [[ds_batch, ds_batch], [ds_batch, ds_batch]]),
+                    },
                 ).invalid_dimensions_for_export(),
             )
             self.assertEqual(
-                {"B": (None, [None, {2: "d=[1]"}, None, {2: "d=[1]"}])},
+                {"B": (None, [[None, {2: "d=[1]"}], [None, {2: "d=[1]"}]])},
                 Cls(
                     (),
                     kwargs,
                     {
                         "A": ds_batch,
-                        "B": (ds_batch, [ds_batch, ds_batch_seq, ds_batch, ds_batch_seq]),
+                        "B": (ds_batch, [[ds_batch, ds_batch_seq], [ds_batch, ds_batch_seq]]),
                     },
                 ).invalid_dimensions_for_export(),
             )

_unittests/ut_helpers/test_cache_helper.py

@@ -7,6 +7,7 @@
 from onnx_diagnostic.torch_export_patches.patch_inputs import (
     convert_dynamic_axes_into_dynamic_shapes,
 )
+from onnx_diagnostic.torch_export_patches import bypass_export_some_errors
 
 
 class TestCacheHelpers(ExtTestCase):
@@ -59,8 +60,9 @@ def test_replace_by(self):
         )
         self.assertEqual(dynamic_shapes, nds)
 
-        cpl = CoupleInputsDynamicShapes(tuple(), kwargs, dynamic_shapes)
-        res = cpl.replace_string_by()
+        with bypass_export_some_errors(patch_transformers=True):
+            cpl = CoupleInputsDynamicShapes(tuple(), kwargs, dynamic_shapes)
+            res = cpl.replace_string_by()
         dsc = res["past_key_values"]
         self.assertEqual([[{0: batch, 2: DYN}], [{0: batch, 2: DYN}]], dsc)
 

_unittests/ut_torch_export_patches/test_patch_inputs.py

@@ -5,6 +5,7 @@
 from onnx_diagnostic.helpers import string_type
 from onnx_diagnostic.torch_export_patches.patch_inputs import (
     convert_dynamic_axes_into_dynamic_shapes,
+    use_dyn_not_str,
 )
 
 
@@ -111,6 +112,26 @@ def test_convert_dynamic_axes_into_dynamic_shapes_2(self):
             string_type(res[1], with_shape=True),
         )
 
+    def test_use_dyn_not_str(self):
+        batch = torch.export.Dim("batch")
+        dynamic_shapes = dict(
+            input_ids={0: batch, 1: "seq"},
+            attention_mask={0: batch, 1: "seq"},
+            position_ids={0: batch, 1: "seq"},
+            past_key_values=[[{0: batch, 2: "seq"}], [{0: batch, 2: "seq"}]],
+        )
+        res = use_dyn_not_str(dynamic_shapes)
+        DYN = torch.export.Dim.DYNAMIC
+        self.assertEqual(
+            dict(
+                input_ids={0: batch, 1: DYN},
+                attention_mask={0: batch, 1: DYN},
+                position_ids={0: batch, 1: DYN},
+                past_key_values=[[{0: batch, 2: DYN}], [{0: batch, 2: DYN}]],
+            ),
+            res,
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/_command_lines_parser.py

@@ -302,7 +302,7 @@ def get_parser_validate() -> ArgumentParser:
 
 def _cmd_validate(argv: List[Any]):
     from .helpers import string_type
-    from .torch_models.test_helper import get_inputs_for_task, validate_model, _ds_clean
+    from .torch_models.test_helper import get_inputs_for_task, validate_model
     from .tasks import supported_tasks
 
     parser = get_parser_validate()
@@ -320,7 +320,7 @@ def _cmd_validate(argv: List[Any]):
             print(f"  + {k.ljust(max_length)}: {string_type(v, with_shape=True)}")
         print("-- dynamic_shapes")
         for k, v in data["dynamic_shapes"].items():
-            print(f"  + {k.ljust(max_length)}: {_ds_clean(v)}")
+            print(f"  + {k.ljust(max_length)}: {string_type(v)}")
     else:
         # Let's skip any invalid combination if known to be unsupported
         if (

onnx_diagnostic/helpers/helper.py

@@ -256,7 +256,7 @@ def string_type(
             print(f"[string_type] L:{type(obj)}")
         return f"{{...}}#{len(obj)}" if with_shape else "{...}"
     # dict
-    if isinstance(obj, dict):
+    if isinstance(obj, dict) and type(obj) is dict:
         if len(obj) == 0:
             if verbose:
                 print(f"[string_type] M:{type(obj)}")
@@ -276,7 +276,7 @@ def string_type(
             )
             for v in obj.values()
         ):
-            # This is dyanmic shapes
+            # This is dynamic shapes
             rows = []
             for k, v in obj.items():
                 if isinstance(v, str):

onnx_diagnostic/torch_export_patches/__init__.py

@@ -2,110 +2,3 @@
     bypass_export_some_errors,
     register_additional_serialization_functions,
 )
-
-"""
--- Missing dependencies --
-
-def is_torchdynamo_exporting() -> bool:
-    "Tells if torch is exporting a model."
-    import torch
-
-    if not hasattr(torch.compiler, "is_exporting"):
-        # torch.compiler.is_exporting requires torch>=2.7
-        return False
-
-    try:
-        return torch.compiler.is_exporting()
-    except Exception:
-        try:
-            import torch._dynamo as dynamo
-
-            return dynamo.is_exporting()  # type: ignore
-        except Exception:
-            return False
-
-
-def string_type(anything, **args):
-    # too long
-    # from onnx_diagnostic.helpers import string_type
-    return str(anything)
-
-
-if pv.Version(transformers.__version__) > pv.Version("4.49.99999"):
-
-    def make_dynamic_cache(
-        key_value_pairs: List[Tuple[torch.Tensor, torch.Tensor]],
-    ) -> transformers.cache_utils.DynamicCache:
-        '''
-        Creates an instance of :class:`transformers.cache_utils.DynamicCache`.
-        This version is valid for ``transformers >= 4.50``.
-
-        :param key_value_pairs: list of pairs of (key, values)
-        :return: :class:`transformers.cache_utils.DynamicCache`
-
-        Example:
-
-        ::
-
-            n_layers = 2
-            bsize, nheads, slen, dim = 2, 4, 3, 7
-
-            past_key_values = make_dynamic_cache(
-                [
-                    (
-                        torch.randn(bsize, nheads, slen, dim),
-                        torch.randn(bsize, nheads, slen, dim),
-                    )
-                    for i in range(n_layers)
-                ]
-            )
-            print(string_type(past_key_values, with_shape=True))
-        '''
-        return transformers.cache_utils.DynamicCache(key_value_pairs)
-
-else:
-
-    def make_dynamic_cache(
-        key_value_pairs: List[Tuple[torch.Tensor, torch.Tensor]],
-    ) -> transformers.cache_utils.DynamicCache:
-        '''
-        Creates an instance of :class:`transformers.cache_utils.DynamicCache`.
-        This version is valid for ``transformers < 4.50``.
-
-        :param key_value_pairs: list of pairs of (key, values)
-        :return: :class:`transformers.cache_utils.DynamicCache`
-
-        Example:
-
-        ::
-
-            n_layers = 2
-            bsize, nheads, slen, dim = 2, 4, 3, 7
-
-            past_key_values = make_dynamic_cache(
-                [
-                    (
-                        torch.randn(bsize, nheads, slen, dim),
-                        torch.randn(bsize, nheads, slen, dim),
-                    )
-                    for i in range(n_layers)
-                ]
-            )
-            print(string_type(past_key_values, with_shape=True))
-        '''
-        cache = transformers.cache_utils.DynamicCache(len(key_value_pairs))
-        for i, (key, value) in enumerate(key_value_pairs):
-            cache.update(key, value, i)
-        return cache
-
-
-def make_encoder_decoder_cache(
-    self_attention_cache: transformers.cache_utils.DynamicCache,
-    cross_attention_cache: transformers.cache_utils.DynamicCache,
-) -> transformers.cache_utils.EncoderDecoderCache:
-    "Creates an EncoderDecoderCache."
-    return transformers.cache_utils.EncoderDecoderCache(
-        self_attention_cache=self_attention_cache, cross_attention_cache=cross_attention_cache
-    )
-
-"""

onnx_diagnostic/torch_export_patches/patch_inputs.py

@@ -183,3 +183,21 @@ def convert_dynamic_axes_into_dynamic_shapes(
             )
 
     return (), updated_kwargs, dynamic_shapes
+
+
+def use_dyn_not_str(dynamic_shapes: Any) -> Any:
+    """
+    Some functions returns dynamic shapes as string.
+    This functions replaces them with ``torch.export.Dim.DYNAMIC``.
+    """
+    if isinstance(dynamic_shapes, list):
+        return [use_dyn_not_str(a) for a in dynamic_shapes]
+    if isinstance(dynamic_shapes, tuple):
+        return tuple(use_dyn_not_str(a) for a in dynamic_shapes)
+    if isinstance(dynamic_shapes, dict):
+        return {k: use_dyn_not_str(v) for k, v in dynamic_shapes.items()}
+    if isinstance(dynamic_shapes, set):
+        return {use_dyn_not_str(a) for a in dynamic_shapes}
+    if isinstance(dynamic_shapes, str):
+        return torch.export.Dim.DYNAMIC
+    return dynamic_shapes