fix ut

Author: xadupre

_unittests/ut_export/test_shape_helper.py

@@ -25,7 +25,6 @@ def test_all_dynamic_shape_from_cache(self):
         ds = all_dynamic_shape_from_inputs(cache)
         self.assertEqual([[{0: "d_0_0", 1: "d_0_1"}], [{0: "d_1_0", 1: "d_1_1"}]], ds)
 
-    @requires_transformers("4.52")
     @requires_torch("2.7.99")
     def test_all_dynamic_shape_all_transformers_cache(self):
         caches = [

_unittests/ut_torch_export_patches/test_patch_serialization_diffusers.py

@@ -49,8 +49,8 @@ def test_unet_2d_condition_output(self):
             # flatten_unflatten
             flat, _spec = torch.utils._pytree.tree_flatten(bo)
             unflat = flatten_unflatten_for_dynamic_shapes(bo, use_dict=True)
-            self.assertIsInstance(unflat, dict)
-            self.assertEqual(list(unflat), ["sample"])
+            self.assertIsInstance(unflat, list)
+            self.assertEqual("#1[T1r3]", self.string_type(unflat))
 
         # export
         class Model(torch.nn.Module):

_unittests/ut_torch_export_patches/test_patch_serialization_transformers.py

@@ -163,8 +163,8 @@ def test_base_model_output_unflatten_flatten(self):
         with torch_export_patches(patch_transformers=True):
             flat, _spec = torch.utils._pytree.tree_flatten(bo)
             unflat = flatten_unflatten_for_dynamic_shapes(bo, use_dict=True)
-            self.assertIsInstance(unflat, dict)
-            self.assertEqual(list(unflat), ["last_hidden_state"])
+            self.assertIsInstance(unflat, list)
+            self.assertEqual("#1[T1r3]", self.string_type(unflat))
 
     @ignore_warnings(UserWarning)
     def test_base_sliding_window_cache_unflatten_flatten(self):
@@ -260,8 +260,10 @@ def test_static_cache(self):
             # flatten_unflatten
             flat, _spec = torch.utils._pytree.tree_flatten(bo)
             unflat = flatten_unflatten_for_dynamic_shapes(bo, use_dict=True)
-            self.assertIsInstance(unflat, dict)
-            self.assertEqual(list(unflat), ["key_cache", "value_cache"])
+            self.assertIsInstance(unflat, list)
+            self.assertEqual(
+                "#2[#3[T1r4,T1r4,T1r4],#3[T1r4,T1r4,T1r4]]", self.string_type(unflat)
+            )
 
         # export
         class Model(torch.nn.Module):

onnx_diagnostic/export/shape_helper.py

@@ -9,7 +9,6 @@ def all_dynamic_shape_from_inputs(inputs: Any, dim_prefix: Any = "d") -> Any:
     All dimensions are considered as dynamic.
     ``dim_prefix`` can be a string (the function uses it as a prefix),
     or ``torch.export.Dim.AUTO`` or ``torch.export.Dim.DYNAMIC``.
-    For caches, ``transformers>=4.52```is better.
 
     .. runpython::
         :showcode: