fix ut

Author: xadupre

_unittests/ut_helpers/test_cache_helper.py

@@ -114,15 +114,21 @@ def test_unflatten_flatten_encoder_decoder_cache(self):
             self.assertEqual(len(unflat[0]), 2)
             self.assertIsInstance(unflat[0][0], list)
             self.assertEqual(len(unflat[0][0]), 3)
+            self.assertEqual(
+                "#2[#3[T1s4x4x4,T1s4x4x4,T1s4x4x4],#3[T1s4x4x4,T1s4x4x4,T1s4x4x4]]",
+                self.string_type(unflat[0], with_shape=True),
+            )
             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]]])",
+                "EncoderDecoderCache(self_attention_cache=DynamicCache("
+                "key_cache=#3[T1s4x4x4,T1s4x4x4,T1s4x4x4], value_cache=#3"
+                "[T1s4x4x4,T1s4x4x4,T1s4x4x4]), cross_attention_cache=DynamicCache"
+                "(key_cache=#3[T1s5x5x5,T1s5x5x5,T1s5x5x5], value_cache=#3"
+                "[T1s5x5x5,T1s5x5x5,T1s5x5x5]))",
                 self.string_type(c2, with_shape=True),
             )
 

_unittests/ut_torch_export_patches/test_patch_inputs.py

@@ -38,7 +38,8 @@ def test_convert_dynamic_axes_into_dynamic_shapes_1(self):
         self.assertEqual(
             (
                 "dict(input_ids:T7s2x8,attention_mask:T7s2x8,position_ids:T7s2x8,"
-                "past_key_values:DynamicCache[serialized](#2[#1[T1s2x1x3x96],#1[T1s2x1x3x96]]))"
+                "past_key_values:DynamicCache(key_cache=#1[T1s2x1x3x96], "
+                "value_cache=#1[T1s2x1x3x96]))"
             ),
             string_type(res[1], with_shape=True),
         )
@@ -107,7 +108,8 @@ def test_convert_dynamic_axes_into_dynamic_shapes_2(self):
         self.assertEqual(
             (
                 "dict(input_ids:T7s2x8,attention_mask:T7s2x8,position_ids:T7s2x8,"
-                "past_key_values:DynamicCache[serialized](#2[#1[T1s2x1x3x96],#1[T1s2x1x3x96]]))"
+                "past_key_values:DynamicCache(key_cache=#1[T1s2x1x3x96], "
+                "value_cache=#1[T1s2x1x3x96]))"
             ),
             string_type(res[1], with_shape=True),
         )

_unittests/ut_torch_export_patches/test_patch_serialization.py

@@ -160,7 +160,7 @@ 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)
+            unflat = flatten_unflatten_for_dynamic_shapes(bo, use_dict=True)
             self.assertIsInstance(unflat, dict)
             self.assertEqual(list(unflat), ["last_hidden_state"])
 

onnx_diagnostic/helpers/cache_helper.py

@@ -18,20 +18,23 @@ def flatten_unflatten_for_dynamic_shapes(obj: Any, use_dict: bool = False) -> An
         for the strict and non strict mode.
     :return: the serialized object
     """
+    if isinstance(obj, torch.Tensor):
+        return obj
     flat, spec = torch.utils._pytree.tree_flatten(obj)
     start = 0
     end = 0
     subtrees = []
     for subspec in spec.children_specs:
         end += subspec.num_leaves
-        if use_dict and subspec.type is dict:
+        if use_dict and (subspec.type is dict or subspec.context):
             value = subspec.unflatten(flat[start:end])
             value = flatten_unflatten_for_dynamic_shapes(value, use_dict=use_dict)
         else:
-            value = flat[start:end]
+            value = subspec.unflatten(flat[start:end])
+            value = flatten_unflatten_for_dynamic_shapes(value, use_dict=use_dict)
         subtrees.append(value)
         start = end
-    if subspec.type is dict:
+    if use_dict and (spec.type is dict or spec.context):
         # This a dictionary.
         return dict(zip(spec.context, subtrees))
     # This is a list.