fix issues

Author: xadupre

_unittests/ut_helpers/test_cache_helper.py

@@ -1,8 +1,13 @@
 import unittest
 import torch
+import transformers
 from onnx_diagnostic.ext_test_case import ExtTestCase
 from onnx_diagnostic.helpers import string_type
-from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
+from onnx_diagnostic.helpers.cache_helper import (
+    make_dynamic_cache,
+    make_encoder_decoder_cache,
+    flatten_unflatten_for_dynamic_shapes,
+)
 from onnx_diagnostic.export import CoupleInputsDynamicShapes
 from onnx_diagnostic.torch_export_patches.patch_inputs import (
     convert_dynamic_axes_into_dynamic_shapes,
@@ -66,6 +71,61 @@ def test_replace_by(self):
         dsc = res["past_key_values"]
         self.assertEqual([[{0: batch, 2: DYN}], [{0: batch, 2: DYN}]], dsc)
 
+    def test_unflatten_flatten_dynamic_cache(self):
+        with bypass_export_some_errors(patch_transformers=True):
+            c1 = make_dynamic_cache([(torch.rand((4, 4, 4)), torch.rand((4, 4, 4)))])
+            self.assertIsInstance(c1, transformers.cache_utils.DynamicCache)
+            unflat = flatten_unflatten_for_dynamic_shapes(c1)
+            self.assertEqual(
+                "#2[#1[T1s4x4x4],#1[T1s4x4x4]]", self.string_type(unflat, with_shape=True)
+            )
+            self.assertEqual(
+                "DynamicCache[serialized](#2[#1[T1s4x4x4],#1[T1s4x4x4]])",
+                self.string_type(c1, with_shape=True),
+            )
+
+    def test_unflatten_flatten_encoder_decoder_cache(self):
+        with bypass_export_some_errors(patch_transformers=True):
+            c2 = make_encoder_decoder_cache(
+                make_dynamic_cache(
+                    [
+                        (torch.rand((4, 4, 4)), torch.rand((4, 4, 4))),
+                        (torch.rand((4, 4, 4)), torch.rand((4, 4, 4))),
+                        (torch.rand((4, 4, 4)), torch.rand((4, 4, 4))),
+                    ]
+                ),
+                make_dynamic_cache(
+                    [
+                        (torch.rand((5, 5, 5)), torch.rand((5, 5, 5))),
+                        (torch.rand((5, 5, 5)), torch.rand((5, 5, 5))),
+                        (torch.rand((5, 5, 5)), torch.rand((5, 5, 5))),
+                    ]
+                ),
+            )
+            self.assertIsInstance(c2, transformers.cache_utils.EncoderDecoderCache)
+            flat, _spec = torch.utils._pytree.tree_flatten(c2)
+            self.assertIsInstance(flat, list)
+            self.assertEqual(len(flat), 12)
+            self.assertIsInstance(flat[0], torch.Tensor)
+            unflat = flatten_unflatten_for_dynamic_shapes(c2)
+            self.assertIsInstance(unflat, list)
+            self.assertEqual(len(unflat), 2)
+            self.assertIsInstance(unflat[0], list)
+            self.assertEqual(len(unflat[0]), 2)
+            self.assertIsInstance(unflat[0][0], list)
+            self.assertEqual(len(unflat[0][0]), 3)
+            self.assertEqual(
+                "#2[#2[#3[T1s4x4x4,T1s4x4x4,T1s4x4x4],#3[T1s4x4x4,T1s4x4x4,T1s4x4x4]],"
+                "#2[#3[T1s5x5x5,T1s5x5x5,T1s5x5x5],#3[T1s5x5x5,T1s5x5x5,T1s5x5x5]]]",
+                self.string_type(unflat, with_shape=True),
+            )
+            self.assertEqual(
+                "EncoderDecoderCache[serialized]("
+                "#2[#2[#3[T1s4x4x4,T1s4x4x4,T1s4x4x4],#3[T1s4x4x4,T1s4x4x4,T1s4x4x4]],"
+                "#2[#3[T1s5x5x5,T1s5x5x5,T1s5x5x5],#3[T1s5x5x5,T1s5x5x5,T1s5x5x5]]])",
+                self.string_type(c2, with_shape=True),
+            )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

_unittests/ut_torch_export_patches/test_patch_serialization.py

@@ -2,7 +2,11 @@
 import torch
 from transformers.modeling_outputs import BaseModelOutput
 from onnx_diagnostic.ext_test_case import ExtTestCase, ignore_warnings
-from onnx_diagnostic.helpers.cache_helper import make_encoder_decoder_cache, make_dynamic_cache
+from onnx_diagnostic.helpers.cache_helper import (
+    make_encoder_decoder_cache,
+    make_dynamic_cache,
+    flatten_unflatten_for_dynamic_shapes,
+)
 from onnx_diagnostic.torch_export_patches.onnx_export_errors import (
     bypass_export_some_errors,
 )
@@ -151,6 +155,15 @@ def forward(self, cache):
         with bypass_export_some_errors():
             torch.export.export(model, (bo,), dynamic_shapes=(ds,))
 
+    @ignore_warnings(UserWarning)
+    def test_base_model_output_unflatten_flatten(self):
+        bo = BaseModelOutput(last_hidden_state=torch.rand((4, 4, 4)))
+        with bypass_export_some_errors(patch_transformers=True):
+            flat, _spec = torch.utils._pytree.tree_flatten(bo)
+            unflat = flatten_unflatten_for_dynamic_shapes(bo)
+            self.assertIsInstance(unflat, dict)
+            self.assertEqual(list(unflat), ["last_hidden_state"])
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/export/dynamic_shapes.py

@@ -316,16 +316,19 @@ def _generic_walker_step(
             return processor(inputs, ds)
         if isinstance(inputs, (int, float, str)):
             return None
-        if isinstance(inputs, (tuple, list, dict)):
-            assert type(ds) is type(
-                inputs
-            ), f"Type mismatch between inputs {type(inputs)} and ds={type(ds)}"
+        if type(inputs) in (tuple, list, dict):
+            # Type must be strict, some custom classes can inherit from those.
+            assert type(inputs) is type(ds), (
+                f"Input type and dynamic shape type mush match but "
+                f"type(inputs)={type(inputs)}, type(ds)={type(ds)}, "
+                f"inputs={string_type(inputs, with_shape=True)}, ds={ds}"
+            )
             assert len(ds) == len(inputs), (
                 f"Length mismatch between inputs {len(inputs)} "
                 f"and ds={len(ds)}\n"
                 f"inputs={string_type(inputs, with_shape=True)}, ds={ds}"
             )
-            if isinstance(inputs, (tuple, list)):
+            if type(inputs) in (tuple, list):
                 value = []
                 for i, d in zip(inputs, ds):
                     value.append(
@@ -338,9 +341,11 @@ def _generic_walker_step(
                     if any(v is not None for v in value)
                     else None
                 )
-            assert set(inputs) == set(
-                ds
-            ), f"Keys mismatch between inputs {set(inputs)} and ds={set(ds)}"
+            assert type(inputs) is dict, f"Unexpected type for inputs {type(inputs)}"
+            assert set(inputs) is set(ds), (
+                f"Keys mismatch between inputs {set(inputs)} and ds={set(ds)}, "
+                f"inputs={string_type(inputs, with_shape=True)}, ds={ds}"
+            )
             dvalue = {}
             for k, v in inputs.items():
                 t = cls._generic_walker_step(

onnx_diagnostic/helpers/cache_helper.py

@@ -19,8 +19,15 @@ def flatten_unflatten_for_dynamic_shapes(obj: Any) -> Any:
     subtrees = []
     for subspec in spec.children_specs:
         end += subspec.num_leaves
-        subtrees.append(subspec.unflatten(flat[start:end]))
+        value = subspec.unflatten(flat[start:end])
+        if not isinstance(value, (torch.Tensor, list)):
+            value = flatten_unflatten_for_dynamic_shapes(value)
+        subtrees.append(value)
         start = end
+    if spec.context:
+        # This a dictionary.
+        return dict(zip(spec.context, subtrees))
+    # This is a list.
     return subtrees
 
 

onnx_diagnostic/tasks/automatic_speech_recognition.py

@@ -74,7 +74,7 @@ def get_inputs(
     shapes = {
         "decoder_input_ids": {0: batch, 1: seq_length},
         "cache_position": {0: seq_length},
-        "encoder_outputs": [{0: batch}],
+        "encoder_outputs": {"last_hidden_state": {0: batch}},
         "past_key_values": [
             [
                 [{0: batch} for _ in range(num_hidden_layers)],