improvment

Author: xadupre

onnx_diagnostic/helpers/helper.py

@@ -252,7 +252,7 @@ def string_type(
         if all(isinstance(k, int) for k in obj):
             return f"{{{s}}}"
         return f"dict({s})"
-    # arrat
+    # array
     if isinstance(obj, np.ndarray):
         from .onnx_helper import np_dtype_to_tensor_dtype
 
@@ -286,6 +286,16 @@ def string_type(
         return "SymInt"
     if isinstance(obj, torch.SymFloat):
         return "SymFloat"
+    if isinstance(obj, torch.export.dynamic_shapes._DimHintType):
+        if obj == torch.export.dynamic_shapes._DimHintType.DYNAMIC:
+            return "DYNAMIC"
+        if obj == torch.export.dynamic_shapes._DimHintType.AUTO:
+            return "AUTO"
+        return str(obj)
+    if obj in (torch.export.Dim.DYNAMIC, torch.export.dynamic_shapes._DimHintType.DYNAMIC):
+        return "DYNAMIC"
+    if obj == (torch.export.Dim.AUTO, torch.export.dynamic_shapes._DimHintType.AUTO):
+        return "AUTO"
     # Tensors
     if isinstance(obj, torch._subclasses.fake_tensor.FakeTensor):
         from .onnx_helper import torch_dtype_to_onnx_dtype
@@ -410,6 +420,10 @@ def string_type(
         )
 
     if obj.__class__.__name__ in ("_DimHint", "_DimHintType"):
+        if obj in (torch.export.Dim.DYNAMIC, torch.export.dynamic_shapes._DimHintType.DYNAMIC):
+            return "DYNAMIC"
+        if obj == (torch.export.Dim.AUTO, torch.export.dynamic_shapes._DimHintType.AUTO):
+            return "AUTO"
         return str(obj)
 
     if isinstance(obj, torch.nn.Module):

onnx_diagnostic/torch_models/test_helper.py

@@ -646,6 +646,18 @@ def call_torch_export_export(
     strict = "nostrict" not in exporter
     args, kwargs = split_args_kwargs(data["inputs_export"])
     ds = data.get("dynamic_shapes", None)
+
+    summary["export_exporter"] = exporter
+    summary["export_optimization"] = optimization or ""
+    summary["export_strict"] = strict
+    summary["export_args"] = string_type(args, with_shape=True)
+    summary["export_kwargs"] = string_type(kwargs, with_shape=True)
+    summary["export_dynamic_shapes"] = string_type(ds)
+
+    # There is an issue with DynamicShape [[],[]] becomes []
+    dse = CoupleInputsDynamicShapes(args, kwargs, ds).replace_string_by()
+    summary["export_dynamic_shapes_export_export"] = string_type(dse)
+
     if verbose:
         print(
             f"[call_torch_export_export] exporter={exporter!r}, "
@@ -654,18 +666,14 @@ def call_torch_export_export(
         print(f"[call_torch_export_export] args={string_type(args, with_shape=True)}")
         print(f"[call_torch_export_export] kwargs={string_type(kwargs, with_shape=True)}")
         print(f"[call_torch_export_export] dynamic_shapes={_ds_clean(ds)}")
+        print(f"[call_torch_export_export] dynamic_shapes_export_export={string_type(dse)}")
         print("[call_torch_export_export] export...")
-    summary["export_exporter"] = exporter
-    summary["export_optimization"] = optimization or ""
-    summary["export_strict"] = strict
-    summary["export_args"] = string_type(args, with_shape=True)
-    summary["export_kwargs"] = string_type(kwargs, with_shape=True)
 
     begin = time.perf_counter()
     if quiet:
         try:
             ep = torch.export.export(
-                data["model"], args, kwargs=kwargs, dynamic_shapes=ds, strict=strict
+                data["model"], args, kwargs=kwargs, dynamic_shapes=dse, strict=strict
             )
         except Exception as e:
             summary["ERR_export_export"] = str(e)
@@ -674,7 +682,7 @@ def call_torch_export_export(
             return summary, data
     else:
         ep = torch.export.export(
-            data["model"], args, kwargs=kwargs, dynamic_shapes=ds, strict=strict
+            data["model"], args, kwargs=kwargs, dynamic_shapes=dse, strict=strict
         )
 
     summary["time_export_export"] = time.perf_counter() - begin