f

Author: xadupre

_unittests/ut_helpers/test_cache_helper.py

@@ -80,7 +80,7 @@ def test_unflatten_flatten_dynamic_cache(self):
                 "#2[#1[T1s4x4x4],#1[T1s4x4x4]]", self.string_type(unflat, with_shape=True)
             )
             self.assertEqual(
-                "DynamicCache[serialized](#2[#1[T1s4x4x4],#1[T1s4x4x4]])",
+                "DynamicCache(key_cache=#1[T1s4x4x4], value_cache=#1[T1s4x4x4])",
                 self.string_type(c1, with_shape=True),
             )
 

_unittests/ut_tasks/try_tasks.py

@@ -184,7 +184,7 @@ def test_automatic_speech_recognition(self):
 
         # generate token ids
         print()
-        with steal_forward(model):
+        with steal_forward(model.model.decoder):
             predicted_ids = model.generate(
                 input_features, forced_decoder_ids=forced_decoder_ids
             )

onnx_diagnostic/export/dynamic_shapes.py

@@ -342,7 +342,7 @@ def _generic_walker_step(
                     else None
                 )
             assert type(inputs) is dict, f"Unexpected type for inputs {type(inputs)}"
-            assert set(inputs) is set(ds), (
+            assert set(inputs) == set(ds), (
                 f"Keys mismatch between inputs {set(inputs)} and ds={set(ds)}, "
                 f"inputs={string_type(inputs, with_shape=True)}, ds={ds}"
             )

onnx_diagnostic/helpers/cache_helper.py

@@ -24,9 +24,11 @@ def flatten_unflatten_for_dynamic_shapes(obj: Any, use_dict: bool = False) -> An
     subtrees = []
     for subspec in spec.children_specs:
         end += subspec.num_leaves
-        value = subspec.unflatten(flat[start:end])
-        if subspec.type is dict:
+        if use_dict and subspec.type is dict:
+            value = subspec.unflatten(flat[start:end])
             value = flatten_unflatten_for_dynamic_shapes(value, use_dict=use_dict)
+        else:
+            value = flat[start:end]
         subtrees.append(value)
         start = end
     if subspec.type is dict: