tiny changes

Author: xadupre

onnx_diagnostic/export/dynamic_shapes.py

@@ -232,8 +232,9 @@ def _generic_walker(
         """
         if not self.args:
             assert isinstance(self.kwargs, dict) and isinstance(self.dynamic_shapes, dict), (
-                f"Type mismatch, args={string_type(self.args)} and "
-                f"dynamic_shapes={self.dynamic_shapes} should have the same type."
+                f"Type mismatch, args={string_type(self.args)}, "
+                f"kwargs={string_type(self.kwargs)} and dynamic_shapes="
+                f"{string_type(self.dynamic_shapes)} should have the same type."
             )
             res = self._generic_walker_step(
                 processor,

onnx_diagnostic/helpers/helper.py

@@ -397,7 +397,7 @@ def string_type(
             return "AUTO"
         if verbose:
             print(f"[string_type] Y7:{type(obj)}")
-        return str(obj)
+        return str(obj).replace("DimHint(DYNAMIC)", "DYNAMIC").replace("DimHint(AUTO)", "AUTO")
 
     if isinstance(obj, bool):
         if with_min_max:
@@ -939,7 +939,7 @@ def flatten_object(x: Any, drop_keys: bool = False) -> Any:
             return flatten_object(list(x.values()), drop_keys=drop_keys)
         return flatten_object(list(x.items()), drop_keys=drop_keys)
 
-    if x.__class__.__name__ in {"DynamicCache", "StaticCache"}:
+    if x.__class__.__name__ in {"DynamicCache", "StaticCache", "HybridCache"}:
         from .cache_helper import CacheKeyValue
 
         kc = CacheKeyValue(x)

onnx_diagnostic/tasks/image_text_to_text.py

@@ -206,8 +206,8 @@ def _check_():
                 batch_size, 1, sequence_length, total_sequence_length
             ),
         ),
-        cache_position=torch.arange(0, sequence_length).to(torch.int64),
         position_ids=torch.arange(0, sequence_length).to(torch.int64).expand((batch_size, -1)),
+        cache_position=torch.arange(0, sequence_length).to(torch.int64),
         past_key_values=make_hybrid_cache(
             [
                 (

onnx_diagnostic/torch_export_patches/eval/model_cases.py

@@ -861,7 +861,7 @@ def forward(self, x):
         y = torch.ones((x.shape[0], dy1))
         return y
 
-    _inputs = [(torch.rand((4, 4)),), (torch.rand((5, 5)),)]
+    _inputs = [(torch.rand((4, 4)),)]
     _dynamic = {"x": {0: DIM("dx"), 1: DIM("dy")}}