fix bug in change_dimnension

Author: xadupre

_unittests/ut_export/test_dynamic_shapes.py

@@ -1,5 +1,6 @@
 import unittest
 import torch
+import transformers
 from onnx_diagnostic.ext_test_case import ExtTestCase, requires_transformers
 from onnx_diagnostic.helpers import string_type
 from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
@@ -742,6 +743,17 @@ def test_couple_input_ds_change_dynamic_dimensions_fixed(self):
         self.assertEqual((1, 5, 8), new_input["A"].shape)
         self.assertEqual((1, 50), new_input["B"].shape)
 
+    def test_couple_input_ds_change_dynamic_dimensions_dynamic_cache(self):
+        inst = CoupleInputsDynamicShapes(
+            (),
+            {"A": make_dynamic_cache([(torch.ones((2, 2, 2, 2)), torch.ones((2, 2, 2, 2)))])},
+            {"A": [[{0: "batch", 2: "last"}], [{0: "batch", 2: "last"}]]},
+        )
+        new_inputs = inst.change_dynamic_dimensions()
+        self.assertIsInstance(new_inputs["A"], transformers.cache_utils.DynamicCache)
+        self.assertEqual((3, 2, 3, 2), new_inputs["A"].key_cache[0].shape)
+        self.assertEqual((3, 2, 3, 2), new_inputs["A"].value_cache[0].shape)
+
     @requires_transformers("4.51")
     def test_dynamic_cache_replace_by_string(self):
         n_layers = 2

_unittests/ut_helpers/test_helper.py

@@ -39,6 +39,8 @@
     dtype_to_tensor_dtype,
 )
 from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache, make_encoder_decoder_cache
+from onnx_diagnostic.torch_models.hghub.hub_api import get_pretrained_config
+
 
 TFLOAT = onnx.TensorProto.FLOAT
 
@@ -484,6 +486,11 @@ def test_flatten_encoder_decoder_cache(self):
         s = string_type(inputs)
         self.assertIn("EncoderDecoderCache", s)
 
+    def test_string_typeçconfig(self):
+        conf = get_pretrained_config("microsoft/phi-2")
+        s = string_type(conf)
+        self.assertStartsWith("PhiConfig(**{", s)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/export/dynamic_shapes.py

@@ -363,16 +363,20 @@ def _generic_walker_step(
         )
         if flatten_unflatten:
             flatunflat = flatten_unflatten_for_dynamic_shapes(inputs)
-            return cls._generic_walker_step(
+            res = cls._generic_walker_step(
                 processor, flatunflat, ds, flatten_unflatten=flatten_unflatten
             )
-        flat, _spec = torch.utils._pytree.tree_flatten(inputs)
+            # Should we restore the original class?
+            return res
+        flat, spec = torch.utils._pytree.tree_flatten(inputs)
         if all(isinstance(t, torch.Tensor) for t in flat):
             # We need to flatten dynamic shapes as well
             ds = flatten_dynamic_shapes(ds)
-        return cls._generic_walker_step(
+        res = cls._generic_walker_step(
             processor, flat, ds, flatten_unflatten=flatten_unflatten
         )
+        # Then we restore the original class.
+        return torch.utils._pytree.tree_unflatten(res, spec)
 
     class ChangeDimensionProcessor:
         def __init__(self, desired_values):

onnx_diagnostic/helpers/helper.py

@@ -666,6 +666,15 @@ def string_type(
             print(f"[string_type] CACHE4:{type(obj)}")
         return f"{obj.__class__.__name__}(...)"
 
+    if obj.__class__.__name__.endswith("Config"):
+        import transformers.configuration_utils as tcu
+
+        if isinstance(obj, tcu.PretrainedConfig):
+            if verbose:
+                print(f"[string_type] CONFIG:{type(obj)}")
+            s = str(obj.to_diff_dict()).replace("\n", "").replace(" ", "")
+            return f"{obj.__class__.__name__}(**{s})"
+
     if verbose:
         print(f"[string_type] END:{type(obj)}")
     raise AssertionError(f"Unsupported type {type(obj).__name__!r} - {type(obj)}")