fix issues

Author: xadupre

_unittests/ut_torch_models/test_hghub_model.py

@@ -112,6 +112,42 @@ def test_get_untrained_model_with_inputs_automatic_speech_recognition(self):
         data = get_untrained_model_with_inputs(mid, verbose=1)
         self.assertIn((data["size"], data["n_weights"]), [(132115968, 33028992)])
         model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        Dim = torch.export.Dim
+        self.maxDiff = None
+        self.assertIn("{0:Dim(batch),1:Dim(seq_length)}", self.string_type(ds))
+        self.assertEqualAny(
+            {
+                "decoder_input_ids": {
+                    0: Dim("batch", min=1, max=1024),
+                    1: Dim("seq_length", min=1, max=4096),
+                },
+                "cache_position": {0: Dim("seq_length", min=1, max=4096)},
+                "encoder_outputs": [{0: Dim("batch", min=1, max=1024)}],
+                "past_key_values": [
+                    [
+                        [
+                            {0: Dim("batch", min=1, max=1024)},
+                            {0: Dim("batch", min=1, max=1024)},
+                        ],
+                        [
+                            {0: Dim("batch", min=1, max=1024)},
+                            {0: Dim("batch", min=1, max=1024)},
+                        ],
+                    ],
+                    [
+                        [
+                            {0: Dim("batch", min=1, max=1024)},
+                            {0: Dim("batch", min=1, max=1024)},
+                        ],
+                        [
+                            {0: Dim("batch", min=1, max=1024)},
+                            {0: Dim("batch", min=1, max=1024)},
+                        ],
+                    ],
+                ],
+            },
+            ds,
+        )
         model(**inputs)
         self.assertEqual(
             "#1[T1r3]",
@@ -125,7 +161,16 @@ def test_get_untrained_model_with_inputs_automatic_speech_recognition(self):
                 "#8[T1r4,T1r4,T1r4,T1r4,T1r4,T1r4,T1r4,T1r4]",
                 self.string_type(flat),
             )
-            torch.export.export(model, (), kwargs=inputs, dynamic_shapes=ds)
+            torch.export.export(model, (), kwargs=inputs, dynamic_shapes=ds, strict=False)
+        with bypass_export_some_errors(patch_transformers=True, verbose=10):
+            flat = torch.utils._pytree.tree_flatten(inputs["past_key_values"])[0]
+            self.assertIsInstance(flat, list)
+            self.assertIsInstance(flat[0], torch.Tensor)
+            self.assertEqual(
+                "#8[T1r4,T1r4,T1r4,T1r4,T1r4,T1r4,T1r4,T1r4]",
+                self.string_type(flat),
+            )
+            torch.export.export(model, (), kwargs=inputs, dynamic_shapes=ds, strict=False)
 
     @hide_stdout()
     def test_get_untrained_model_with_inputs_imagetext2text_generation(self):

onnx_diagnostic/ext_test_case.py

@@ -910,6 +910,9 @@ def assertEqualAny(
         elif hasattr(expected, "shape"):
             self.assertEqual(type(expected), type(value), msg=msg)
             self.assertEqualArray(expected, value, msg=msg, atol=atol, rtol=rtol)
+        elif expected.__class__.__name__ in ("Dim", "_Dim"):
+            self.assertEqual(type(expected), type(value), msg=msg)
+            self.assertEqual(expected.__name__, value.__name__, msg=msg)
         else:
             raise AssertionError(
                 f"Comparison not implemented for types {type(expected)} and {type(value)}"

onnx_diagnostic/helpers/helper.py

@@ -249,6 +249,8 @@ def string_type(
             limit=limit,
         )
         s = ",".join(f"{kv[0]}:{string_type(kv[1],**kws)}" for kv in obj.items())
+        if all(isinstance(k, int) for k in obj):
+            return f"{{{s}}}"
         return f"dict({s})"
     # arrat
     if isinstance(obj, np.ndarray):
@@ -279,7 +281,7 @@ def string_type(
     if isinstance(obj, torch.export.dynamic_shapes._DerivedDim):
         return "DerivedDim"
     if isinstance(obj, torch.export.dynamic_shapes._Dim):
-        return "Dim"
+        return f"Dim({obj.__name__})"
     if isinstance(obj, torch.SymInt):
         return "SymInt"
     if isinstance(obj, torch.SymFloat):
@@ -355,6 +357,11 @@ def string_type(
     if isinstance(obj, slice):
         return "slice"
 
+    if obj == torch.export.Dim.DYNAMIC:
+        return "DYNAMIC"
+    if obj == torch.export.Dim.AUTO:
+        return "AUTO"
+
     # others classes
 
     if obj.__class__ in torch.utils._pytree.SUPPORTED_NODES:

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -36,7 +36,7 @@ def patch_module_or_classes(mod, verbose: int = 0) -> Dict[type, Dict[type, Call
         original = cls._PATCHED_CLASS_
         methods = cls._PATCHES_
         if verbose:
-            print(f"[patch_module_or_classes] {name} - {cls.__name__}: {', '.join(methods)}")
+            print(f"[patch_module_or_classes] {name}.{cls.__name__}: {', '.join(methods)}")
 
         keep = {n: getattr(original, n, None) for n in methods}
         for n in methods:
@@ -69,7 +69,7 @@ def unpatch_module_or_classes(mod, info: Dict[type, Dict[type, Callable]], verbo
     for cls, methods in info.items():
         assert cls in set_patch, f"No patch registered for {cls} in {mod} (found {set_patch})"
         if verbose:
-            print(f"[unpatch_module_or_classes] {name} - {cls.__name__}: {', '.join(methods)}")
+            print(f"[unpatch_module_or_classes] {name}.{cls.__name__}: {', '.join(methods)}")
         original = cls._PATCHED_CLASS_
         for n, v in methods.items():
             if v is None:

onnx_diagnostic/torch_export_patches/onnx_export_serialization.py

@@ -1,16 +1,58 @@
 import pprint
-from typing import Any, Dict, List, Set, Tuple
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple
 import packaging.version as pv
 import optree
 import torch
 import transformers
 from transformers.cache_utils import DynamicCache, MambaCache, EncoderDecoderCache
 from transformers.modeling_outputs import BaseModelOutput
+from ..helpers import string_type
 
 
 PATCH_OF_PATCHES: Set[Any] = set()
 
 
+def _register_class_serialization(
+    cls,
+    f_flatten: Callable,
+    f_unflatten: Callable,
+    f_flatten_with_keys: Callable,
+    f_check: Optional[Callable] = None,
+    verbose: int = 0,
+) -> bool:
+    if cls is not None and cls in torch.utils._pytree.SUPPORTED_NODES:
+        return False
+
+    if verbose:
+        print(f"[_register_cache_serialization] register {cls}")
+    torch.utils._pytree.register_pytree_node(
+        cls,
+        f_flatten,
+        f_unflatten,
+        serialized_type_name=f"{cls.__module__}.{cls.__name__}",
+        flatten_with_keys_fn=f_flatten_with_keys,
+    )
+    if pv.Version(torch.__version__) < pv.Version("2.7"):
+        if verbose:
+            print(
+                f"[_register_cache_serialization] "
+                f"register {cls} for torch=={torch.__version__}"
+            )
+        torch.fx._pytree.register_pytree_flatten_spec(cls, lambda x, _: f_flatten(x)[0])
+
+    # check
+    if f_check:
+        inst = f_check()
+        values, spec = torch.utils._pytree.tree_flatten(inst)
+        restored = torch.utils._pytree.tree_unflatten(values, spec)
+        assert string_type(inst, with_shape=True) == string_type(restored, with_shape=True), (
+            f"Issue with registration of class {cls} "
+            f"inst={string_type(inst, with_shape=True)}, "
+            f"restored={string_type(restored, with_shape=True)}"
+        )
+    return True
+
+
 def _register_cache_serialization(verbose: int = 0) -> Dict[str, bool]:
     # DynamicCache serialization is different in transformers and does not
     # play way with torch.export.export.
@@ -28,26 +70,20 @@ def _register_cache_serialization(verbose: int = 0) -> Dict[str, bool]:
     ):
         if verbose:
             print(
-                "[_register_cache_serialization] DynamicCache "
-                "is unregistered and registered first."
+                f"[_fix_registration] DynamicCache is unregistered and "
+                f"registered first for transformers=={transformers.__version__}"
             )
         _unregister(DynamicCache, verbose=verbose)
-        torch.utils._pytree.register_pytree_node(
+        _register_class_serialization(
             DynamicCache,
             flatten_dynamic_cache,
             unflatten_dynamic_cache,
-            serialized_type_name=f"{DynamicCache.__module__}.{DynamicCache.__name__}",
-            flatten_with_keys_fn=flatten_with_keys_dynamic_cache,
+            flatten_with_keys_dynamic_cache,
+            # f_check=make_dynamic_cache([(torch.rand((4, 4, 4)), torch.rand((4, 4, 4)))]),
+            verbose=verbose,
         )
         if verbose:
-            print(
-                "[_register_cache_serialization] DynamicCache "
-                "unregistered and registered done."
-            )
-        if pv.Version(torch.__version__) < pv.Version("2.7"):
-            torch.fx._pytree.register_pytree_flatten_spec(
-                DynamicCache, lambda x, _: [x.key_cache, x.value_cache]
-            )
+            print("[_fix_registration] DynamicCache done.")
         # To avoid doing it multiple times.
         PATCH_OF_PATCHES.add(DynamicCache)
 
@@ -59,120 +95,52 @@ def _register_cache_serialization(verbose: int = 0) -> Dict[str, bool]:
     ):
         if verbose:
             print(
-                "[_register_cache_serialization] BaseModelOutput "
-                "is unregistered and registered first."
+                f"[_fix_registration] BaseModelOutput is unregistered and "
+                f"registered first for transformers=={transformers.__version__}"
             )
         _unregister(BaseModelOutput, verbose=verbose)
-        torch.utils._pytree.register_pytree_node(
+        _register_class_serialization(
             BaseModelOutput,
             flatten_base_model_output,
             unflatten_base_model_output,
-            serialized_type_name=f"{BaseModelOutput.__module__}.{BaseModelOutput.__name__}",
-            flatten_with_keys_fn=flatten_with_keys_base_model_output,
+            flatten_with_keys_base_model_output,
+            verbose=verbose,
         )
         if verbose:
-            print(
-                "[_register_cache_serialization] BaseModelOutput "
-                "unregistered and registered done."
-            )
+            print("[_fix_registration] BaseModelOutput done.")
 
         # To avoid doing it multiple times.
         PATCH_OF_PATCHES.add(BaseModelOutput)
 
-    unregistered_dynamic_cache = True
-    if DynamicCache is not None and DynamicCache in torch.utils._pytree.SUPPORTED_NODES:
-        if verbose > 1:
-            print(f"[_register_cache_serialization] {DynamicCache} already registered")
-        unregistered_dynamic_cache = False
-    else:
-        if verbose:
-            print("[_register_cache_serialization] register DynamicCache")
-        torch.utils._pytree.register_pytree_node(
-            DynamicCache,
-            flatten_dynamic_cache,
-            unflatten_dynamic_cache,
-            serialized_type_name=f"{DynamicCache.__module__}.{DynamicCache.__name__}",
-            flatten_with_keys_fn=flatten_with_keys_dynamic_cache,
-        )
-        if pv.Version(torch.__version__) < pv.Version("2.7"):
-            torch.fx._pytree.register_pytree_flatten_spec(
-                DynamicCache, lambda x, _: [x.key_cache, x.value_cache]
-            )
-
-        # check
-        from ..helpers.cache_helper import make_dynamic_cache
-
-        cache = make_dynamic_cache([(torch.rand((4, 4, 4)), torch.rand((4, 4, 4)))])
-        values, spec = torch.utils._pytree.tree_flatten(cache)
-        cache2 = torch.utils._pytree.tree_unflatten(values, spec)
-        # torch.fx._pytree.tree_flatten(cache)
-        assert len(cache2.key_cache) == 1
-
-    # BaseModelOutput
-    unregistered_base_model_output = True
-    if BaseModelOutput is not None and BaseModelOutput in torch.utils._pytree.SUPPORTED_NODES:
-        if verbose > 1:
-            print(f"[_register_cache_serialization] {BaseModelOutput} already registered")
-        # It is already registered because bypass_export_some_errors was called
-        # within a section already calling bypass_export_some_errors or transformers
-        # has updated its code to do it.
-        # No need to register and unregister then.
-        unregistered_base_model_output = False
-    else:
-        if verbose:
-            print("[_register_cache_serialization] register BaseModelOutput")
-        torch.utils._pytree.register_pytree_node(
-            BaseModelOutput,
-            flatten_encoder_decoder_cache,
-            unflatten_encoder_decoder_cache,
-            serialized_type_name=f"{BaseModelOutput.__module__}.{BaseModelOutput.__name__}",
-            flatten_with_keys_fn=flatten_with_keys_base_model_output,
-        )
-
-    # MambaCache
-    unregistered_mamba_cache = True
-    if MambaCache in torch.utils._pytree.SUPPORTED_NODES:
-        if verbose > 1:
-            print(f"[_register_cache_serialization] {MambaCache} already registered")
-        # It is already registered because bypass_export_some_errors was called
-        # within a section already calling bypass_export_some_errors or transformers
-        # has updated its code to do it.
-        # No need to register and unregister then.
-        unregistered_mamba_cache = False
-    else:
-        if verbose:
-            print("[_register_cache_serialization] register MambaCache")
-        torch.utils._pytree.register_pytree_node(
-            MambaCache,
-            flatten_mamba_cache,
-            unflatten_mamba_cache,
-            serialized_type_name=f"{MambaCache.__module__}.{MambaCache.__name__}",
-            flatten_with_keys_fn=flatten_with_keys_mamba_cache,
-        )
-
-    # EncoderDecoderCache
-    unregistered_encode_decode_cache = True
-    if (
-        EncoderDecoderCache is not None
-        and EncoderDecoderCache in torch.utils._pytree.SUPPORTED_NODES
-    ):
-        if verbose > 1:
-            print(f"[_register_cache_serialization] {EncoderDecoderCache} already registered")
-        # It is already registered because bypass_export_some_errors was called
-        # within a section already calling bypass_export_some_errors or transformers
-        # has updated its code to do it.
-        # No need to register and unregister then.
-        unregistered_encode_decode_cache = False
-    else:
-        if verbose:
-            print("[_register_cache_serialization] register EncoderDecoderCache")
-        torch.utils._pytree.register_pytree_node(
-            EncoderDecoderCache,
-            flatten_encoder_decoder_cache,
-            unflatten_encoder_decoder_cache,
-            serialized_type_name=f"{EncoderDecoderCache.__module__}.{EncoderDecoderCache.__name__}",
-            flatten_with_keys_fn=flatten_with_keys_encoder_decoder_cache,
-        )
+    unregistered_dynamic_cache = _register_class_serialization(
+        DynamicCache,
+        flatten_dynamic_cache,
+        unflatten_dynamic_cache,
+        flatten_with_keys_dynamic_cache,
+        # f_check=make_dynamic_cache([(torch.rand((4, 4, 4)), torch.rand((4, 4, 4)))]),
+        verbose=verbose,
+    )
+    unregistered_base_model_output = _register_class_serialization(
+        BaseModelOutput,
+        flatten_base_model_output,
+        unflatten_base_model_output,
+        flatten_with_keys_base_model_output,
+        verbose=verbose,
+    )
+    unregistered_encode_decode_cache = _register_class_serialization(
+        EncoderDecoderCache,
+        flatten_encoder_decoder_cache,
+        unflatten_encoder_decoder_cache,
+        flatten_with_keys_encoder_decoder_cache,
+        verbose=verbose,
+    )
+    unregistered_mamba_cache = _register_class_serialization(
+        MambaCache,
+        flatten_mamba_cache,
+        unflatten_mamba_cache,
+        flatten_with_keys_mamba_cache,
+        verbose=verbose,
+    )
 
     return dict(
         DynamicCache=unregistered_dynamic_cache,
@@ -213,8 +181,6 @@ def _unregister_cache_serialization(undo: Dict[str, bool], verbose: int = 0):
     for cls in [MambaCache, DynamicCache, EncoderDecoderCache, BaseModelOutput]:
         if undo.get(cls.__name__, False):
             _unregister(cls, verbose)
-        elif verbose > 1:
-            print(f"[_unregister_cache_serialization] skip unregister {cls.__name__}")
 
 
 ############

onnx_diagnostic/torch_models/test_helper.py

@@ -23,16 +23,7 @@ def empty(value: Any) -> bool:
 
 
 def _ds_clean(v):
-    return (
-        str(v)
-        .replace(",min=None", "")
-        .replace(",max=None", "")
-        .replace(",_factory=True", "")
-        .replace("<class 'onnx_diagnostic.torch_models.hghub.model_inputs.", "")
-        .replace("'>", "")
-        .replace("_DimHint(type=<_DimHintType.DYNAMIC: 3>)", "DYNAMIC")
-        .replace("_DimHint(type=<_DimHintType.AUTO: 3>)", "AUTO")
-    )
+    return string_type(v)
 
 
 def get_inputs_for_task(task: str, config: Optional[Any] = None) -> Dict[str, Any]: