fix issues

Author: xadupre

_doc/recipes/plot_dynamic_shapes_json.py

@@ -75,11 +75,21 @@ def flatten_unflatten_like_dynamic_shapes(obj):
         value = flatten_unflatten_like_dynamic_shapes(value)
         subtrees.append(value)
         start = end
-    if spec.type is dict or spec.context:
+    if spec.type is dict:
+        # This a dictionary.
         return dict(zip(spec.context, subtrees))
     if spec.type is tuple:
         return tuple(subtrees)
-    return subtrees
+    if spec.type is list:
+        return list(subtrees)
+    if spec.context:
+        # This is a custom class with attributes.
+        # It is returned as a list.
+        return list(subtrees)
+    raise ValueError(
+        f"Unable to interpret spec type {spec.type} "
+        f"(type is {type(spec.type)}, context is {spec.context})."
+    )
 
 
 def _align(inputs, ds):

_unittests/ut_export/test_shape_helper.py

@@ -10,7 +10,6 @@
     make_sliding_window_cache,
     make_encoder_decoder_cache,
     make_static_cache,
-    make_mamba_cache,
 )
 from onnx_diagnostic.torch_models.hghub import get_untrained_model_with_inputs
 from onnx_diagnostic.torch_export_patches import torch_export_patches
@@ -119,27 +118,6 @@ def test_all_dynamic_shape_all_transformers_cache(self):
                     ],
                 ],
             ),
-            (
-                make_mamba_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))),
-                    ]
-                ),
-                [
-                    [
-                        {0: "d_0_0", 1: "d_0_1", 2: "d_0_2"},
-                        {0: "d_1_0", 1: "d_1_1", 2: "d_1_2"},
-                        {0: "d_2_0", 1: "d_2_1", 2: "d_2_2"},
-                    ],
-                    [
-                        {0: "d_3_0", 1: "d_3_1", 2: "d_3_2"},
-                        {0: "d_4_0", 1: "d_4_1", 2: "d_4_2"},
-                        {0: "d_5_0", 1: "d_5_1", 2: "d_5_2"},
-                    ],
-                ],
-            ),
         ]
         with torch_export_patches(patch_transformers=True):
             for cache, exds in caches: