strict=False by default

Author: xadupre

onnx_diagnostic/helpers/cache_helper.py

@@ -5,12 +5,17 @@
 import transformers.cache_utils
 
 
-def flatten_unflatten_for_dynamic_shapes(obj: Any) -> Any:
+def flatten_unflatten_for_dynamic_shapes(obj: Any, use_dict: bool = False) -> Any:
     """
     Returns the object in a different structure similar to what
     the definition of the dynamic shapes should use.
 
     :param obj: object from a custom class
+    :param use_dict: closer to the original result but
+        :func:`torch.export.export` only considers the values,
+        the context gives the dictionary keys but it is not expressed
+        in the dynamic shapes, these specifications seems to be different
+        for the strict and non strict mode.
     :return: the serialized object
     """
     flat, spec = torch.utils._pytree.tree_flatten(obj)
@@ -20,11 +25,11 @@ def flatten_unflatten_for_dynamic_shapes(obj: Any) -> Any:
     for subspec in spec.children_specs:
         end += subspec.num_leaves
         value = subspec.unflatten(flat[start:end])
-        if not isinstance(value, (torch.Tensor, list)):
-            value = flatten_unflatten_for_dynamic_shapes(value)
+        if subspec.type is dict:
+            value = flatten_unflatten_for_dynamic_shapes(value, use_dict=use_dict)
         subtrees.append(value)
         start = end
-    if spec.context:
+    if subspec.type is dict:
         # This a dictionary.
         return dict(zip(spec.context, subtrees))
     # This is a list.

onnx_diagnostic/helpers/helper.py

@@ -513,6 +513,72 @@ def string_type(
 
     # others classes
 
+    if obj.__class__.__name__ == "MambaCache":
+        c = string_type(
+            obj.conv_states,
+            with_shape=with_shape,
+            with_min_max=with_min_max,
+            with_device=with_device,
+            limit=limit,
+            verbose=verbose,
+        )
+        d = string_type(
+            obj.ssm_states,
+            with_shape=with_shape,
+            with_min_max=with_min_max,
+            with_device=with_device,
+            limit=limit,
+            verbose=verbose,
+        )
+        if verbose:
+            print(f"[string_type] CACHE1:{type(obj)}")
+        return f"MambaCache(conv_states={c}, ssm_states={d})"
+
+    if obj.__class__.__name__ == "DynamicCache":
+        kc = string_type(
+            obj.key_cache,
+            with_shape=with_shape,
+            with_min_max=with_min_max,
+            with_device=with_device,
+            limit=limit,
+            verbose=verbose,
+        )
+        vc = string_type(
+            obj.value_cache,
+            with_shape=with_shape,
+            with_min_max=with_min_max,
+            with_device=with_device,
+            limit=limit,
+            verbose=verbose,
+        )
+        if verbose:
+            print(f"[string_type] CACHE2:{type(obj)}")
+        return f"{obj.__class__.__name__}(key_cache={kc}, value_cache={vc})"
+
+    if obj.__class__.__name__ == "EncoderDecoderCache":
+        att = string_type(
+            obj.self_attention_cache,
+            with_shape=with_shape,
+            with_min_max=with_min_max,
+            with_device=with_device,
+            limit=limit,
+            verbose=verbose,
+        )
+        cross = string_type(
+            obj.cross_attention_cache,
+            with_shape=with_shape,
+            with_min_max=with_min_max,
+            with_device=with_device,
+            limit=limit,
+            verbose=verbose,
+        )
+        if verbose:
+            print(f"[string_type] CACHE3:{type(obj)}")
+        return (
+            f"{obj.__class__.__name__}(self_attention_cache={att}, "
+            f"cross_attention_cache={cross})"
+        )
+
     if obj.__class__ in torch.utils._pytree.SUPPORTED_NODES:
         from .cache_helper import flatten_unflatten_for_dynamic_shapes
 
@@ -595,74 +661,6 @@ def string_type(
             print(f"[string_type] TT8:{type(obj)}")
         return repr(obj).replace(" ", "").replace("\n", " ")
 
-    # to avoid failures
-
-    if obj.__class__.__name__ == "MambaCache":
-        c = string_type(
-            obj.conv_states,
-            with_shape=with_shape,
-            with_min_max=with_min_max,
-            with_device=with_device,
-            limit=limit,
-            verbose=verbose,
-        )
-        d = string_type(
-            obj.ssm_states,
-            with_shape=with_shape,
-            with_min_max=with_min_max,
-            with_device=with_device,
-            limit=limit,
-            verbose=verbose,
-        )
-        if verbose:
-            print(f"[string_type] CACHE1:{type(obj)}")
-        return f"MambaCache(conv_states={c}, ssm_states={d})"
-
-    if obj.__class__.__name__ == "DynamicCache":
-        kc = string_type(
-            obj.key_cache,
-            with_shape=with_shape,
-            with_min_max=with_min_max,
-            with_device=with_device,
-            limit=limit,
-            verbose=verbose,
-        )
-        vc = string_type(
-            obj.value_cache,
-            with_shape=with_shape,
-            with_min_max=with_min_max,
-            with_device=with_device,
-            limit=limit,
-            verbose=verbose,
-        )
-        if verbose:
-            print(f"[string_type] CACHE2:{type(obj)}")
-        return f"{obj.__class__.__name__}(key_cache={kc}, value_cache={vc})"
-
-    if obj.__class__.__name__ == "EncoderDecoderCache":
-        att = string_type(
-            obj.self_attention_cache,
-            with_shape=with_shape,
-            with_min_max=with_min_max,
-            with_device=with_device,
-            limit=limit,
-            verbose=verbose,
-        )
-        cross = string_type(
-            obj.cross_attention_cache,
-            with_shape=with_shape,
-            with_min_max=with_min_max,
-            with_device=with_device,
-            limit=limit,
-            verbose=verbose,
-        )
-        if verbose:
-            print(f"[string_type] CACHE3:{type(obj)}")
-        return (
-            f"{obj.__class__.__name__}(self_attention_cache={att}, "
-            f"cross_attention_cache={cross})"
-        )
-
     if ignore:
         if verbose:
             print(f"[string_type] CACHE4:{type(obj)}")

onnx_diagnostic/tasks/automatic_speech_recognition.py

@@ -69,12 +69,12 @@ def get_inputs(
         )
     """
     batch = torch.export.Dim("batch", min=1, max=1024)
-    seq_length = torch.export.Dim("seq_length", min=1, max=4096)
+    seq_length = "seq_length"
 
     shapes = {
         "decoder_input_ids": {0: batch, 1: seq_length},
         "cache_position": {0: seq_length},
-        "encoder_outputs": {"last_hidden_state": {0: batch}},
+        "encoder_outputs": [{0: batch}],  # last_hidden_state
         "past_key_values": [
             [
                 [{0: batch} for _ in range(num_hidden_layers)],

onnx_diagnostic/torch_models/test_helper.py

@@ -12,6 +12,7 @@
 from ..helpers.torch_test_helper import to_any, torch_deepcopy
 from ..helpers.cache_helper import flatten_unflatten_for_dynamic_shapes
 from ..torch_export_patches import bypass_export_some_errors
+from ..torch_export_patches.patch_inputs import use_dyn_not_str
 from .hghub import get_untrained_model_with_inputs
 from .hghub.model_inputs import random_input_kwargs
 
@@ -633,14 +634,15 @@ def call_torch_export_export(
         another one with whatever the function produces
     """
     assert exporter in {
+        "export",
         "export-strict",
         "export-nostrict",
     }, f"Unexpected value for exporter={exporter!r}"
     assert not optimization, f"No optimization is implemented for exporter={exporter!r}"
     assert "model" in data, f"model is missing from data: {sorted(data)}"
     assert "inputs_export" in data, f"inputs_export is missing from data: {sorted(data)}"
     summary: Dict[str, Union[str, int, float]] = {}
-    strict = "nostrict" not in exporter
+    strict = "-strict" in exporter
     args, kwargs = split_args_kwargs(data["inputs_export"])
     ds = data.get("dynamic_shapes", None)
 
@@ -652,7 +654,9 @@ def call_torch_export_export(
     summary["export_dynamic_shapes"] = string_type(ds)
 
     # There is an issue with DynamicShape [[],[]] becomes []
-    dse = CoupleInputsDynamicShapes(args, kwargs, ds).replace_string_by()
+    dse = use_dyn_not_str(ds)
+    # dse = CoupleInputsDynamicShapes(args, kwargs, ds).replace_string_by()
+
     summary["export_dynamic_shapes_export_export"] = string_type(dse)
 
     if verbose:
@@ -1015,7 +1019,7 @@ def call_torch_export_custom(
     assert "model" in data, f"model is missing from data: {sorted(data)}"
     assert "inputs_export" in data, f"inputs_export is missing from data: {sorted(data)}"
     summary: Dict[str, Union[str, int, float]] = {}
-    dynamo = "nostrict" not in exporter
+    strict = "-strict" in exporter
     args, kwargs = split_args_kwargs(data["inputs_export"])
     ds = data.get("dynamic_shapes", None)
     if verbose:
@@ -1029,15 +1033,15 @@ def call_torch_export_custom(
         print("[call_torch_export_custom] export...")
     summary["export_exporter"] = exporter
     summary["export_optimization"] = optimization or ""
-    summary["export_dynamo"] = dynamo
+    summary["export_strict"] = strict
     summary["export_args"] = string_type(args, with_shape=True)
     summary["export_kwargs"] = string_type(kwargs, with_shape=True)
 
     from experimental_experiment.torch_interpreter import to_onnx, ExportOptions
     from experimental_experiment.xbuilder import OptimizationOptions
 
     export_options = ExportOptions(
-        strict="nostrict" not in exporter,
+        strict=strict,
         decomposition_table=(
             "dec" if "-dec" in exporter else ("all" if "-all" in exporter else None)
         ),
@@ -1057,6 +1061,7 @@ def call_torch_export_custom(
                 optimize=bool(optimization),
                 large_model=True,
                 return_optimize_report=True,
+                verbose=max(verbose - 2, 0),
             )
         except Exception as e:
             summary["ERR_export_export"] = str(e)
@@ -1074,6 +1079,7 @@ def call_torch_export_custom(
             optimize=bool(optimization),
             large_model=True,
             return_optimize_report=True,
+            verbose=max(verbose - 2, 0),
         )
 
     new_stat = {}