fix fake tensors

Author: xadupre

_unittests/ut_torch_models/test_validate_whole_models1.py

@@ -226,6 +226,7 @@ def test_n_validate_phi35_mini_instruct(self):
         self.clean_dump()
 
     @hide_stdout()
+    @requires_transformers("4.57")
     def test_o_validate_model_export_fake(self):
         mid = "arnir0/Tiny-LLM"
         summary, data = validate_model(

onnx_diagnostic/export/shape_helper.py

@@ -210,6 +210,7 @@ def make_fake_with_dynamic_dimensions(
     This uses function :func:`onnx_diagnostic.helpers.fake_tensor_helper.make_fake`.
     Parameter ``existing`` is used to reused the same object when the dynamic
     dimension is given the same name as another one.
+    This function works with caches only if ``transformers>=4.57``.
 
     A simple tensor:
 

onnx_diagnostic/helpers/fake_tensor_helper.py

@@ -144,17 +144,18 @@ def make_fake_with_dynamic_dimensions(self, x: Any, dynamic_shapes: Any) -> Any:
         """
         See
         :func:`onnx_diagnostic.export.shape_helper.make_fake_with_dynamic_dimensions`.
+        If caches are used, it requires ``transformers>=4.57``.
         """
         if x is None:
             return None, None
-        if isinstance(x, (list, tuple)):
+        if type(x) in (list, tuple):
             return x.__class__(
                 [
                     self.make_fake_with_dynamic_dimensions(i, dynamic_shapes=ds)
                     for i, ds in zip(x, dynamic_shapes)
                 ]
             )
-        if isinstance(x, dict):
+        if type(x) is dict:
             return {
                 k: self.make_fake_with_dynamic_dimensions(v, dynamic_shapes=dynamic_shapes[k])
                 for k, v in x.items()
@@ -187,6 +188,17 @@ def make_fake_with_dynamic_dimensions(self, x: Any, dynamic_shapes: Any) -> Any:
                 x.cross_attention_cache, dynamic_shapes=dynamic_shapes[1]
             )
             return x
+        if x.__class__.__name__ == "BaseModelOutput":
+            assert (
+                list(x.keys()) == ["last_hidden_state"] and x.last_hidden_state is not None
+            ), (
+                f"Field 'last_hidden_state' is empty for {type(x)} or other fields "
+                f"{list(x.keys())} are used."
+            )
+            x.last_hidden_state = self.make_fake_with_dynamic_dimensions(
+                x.last_hidden_state, dynamic_shapes=dynamic_shapes[0]
+            )
+            return x
         if hasattr(x, "shape"):
             assert dynamic_shapes is None or isinstance(dynamic_shapes, dict), (
                 f"dynamic_shapes must be a dictionary at this stage but "