onnx-diagnostic 0.7.12__py3-none-any.whl → 0.7.13__py3-none-any.whl

onnx_diagnostic/__init__.py

@@ -3,5 +3,5 @@ Patches, Investigates onnx models.
 Functions, classes to dig into a model when this one is right, slow, wrong...
 """
 
-__version__ = "0.7.12"
+__version__ = "0.7.13"
 __author__ = "Xavier Dupré"

onnx_diagnostic/export/dynamic_shapes.py

@@ -56,6 +56,14 @@ class CoupleInputsDynamicShapes:
         self.kwargs = kwargs
         self.dynamic_shapes = dynamic_shapes
         self.args_names = args_names
+        if not self.kwargs and isinstance(self.dynamic_shapes, dict):
+            # This assumes the dictionary for the dynamic shapes is ordered
+            # the same way the args are. The input names are not known.
+            assert len(self.dynamic_shapes) == len(self.args), (
+                f"Length mismatch, kwargs is empty, len(dynamic_shapes)="
+                f"{len(self.dynamic_shapes)}, len(args)={len(self.args)}"
+            )
+            self.dynamic_shapes = tuple(self.dynamic_shapes.values())
 
     def __str__(self) -> str:
         return "\n".join(
@@ -232,8 +240,9 @@ class CoupleInputsDynamicShapes:
         """
         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:
@@ -516,8 +516,10 @@ def string_type(
                 print(f"[string_type] V2:{type(obj)}")
             return "OV(NOTENSOR)"
         if with_min_max:
+            from .torch_helper import to_numpy
+
             try:
-                t = obj.numpy()
+                t = to_numpy(obj)
             except Exception:
                 # pass unable to convert into numpy (bfloat16, ...)
                 if verbose:
@@ -939,7 +941,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)
@@ -1233,9 +1235,13 @@ def max_diff(
 
     if isinstance(expected, np.ndarray) or isinstance(got, np.ndarray):
         if isinstance(expected, torch.Tensor):
-            expected = expected.detach().cpu().numpy()
+            from .torch_helper import to_numpy
+
+            expected = to_numpy(expected)
         if isinstance(got, torch.Tensor):
-            got = got.detach().cpu().numpy()
+            from .torch_helper import to_numpy
+
+            got = to_numpy(got)
         if verbose >= 6:
             print(f"[max_diff] tensor: {string_type(expected)} ? {string_type(got)}")
 

onnx_diagnostic/helpers/mini_onnx_builder.py

@@ -381,6 +381,23 @@ def _flatten_iterator(obj: Any, sep: str) -> Iterator:
                 else:
                     for p, o in _flatten_iterator(getattr(obj, att), sep):
                         yield f"DynamicCache_{att}{sep}{p}", o
+        elif obj.__class__.__name__ == "StaticCache":
+            # transformers
+            import transformers
+            from .cache_helper import CacheKeyValue
+
+            assert isinstance(
+                obj, transformers.cache_utils.StaticCache
+            ), f"Unexpected type {type(obj)}"
+            obj = CacheKeyValue(obj)
+            atts = ["key_cache", "value_cache"]
+            for i, att in enumerate(atts):
+                if i == len(atts) - 1:
+                    for p, o in _flatten_iterator(getattr(obj, att), sep):
+                        yield f"StaticCache._{att}{sep}{p}", o
+                else:
+                    for p, o in _flatten_iterator(getattr(obj, att), sep):
+                        yield f"StaticCache_{att}{sep}{p}", o
         else:
             raise NotImplementedError(f"Unexpected type {type(obj)}")
 

onnx_diagnostic/helpers/model_builder_helper.py

@@ -203,6 +203,7 @@ def create_model_builder(
         "ChatGLMModel": builder.ChatGLMModel,
         "Ernie4_5_ForCausalLM": builder.ErnieModel,
         "GemmaForCausalLM": builder.Gemma2Model,
+        "Gemma2ForCausalLM": builder.Gemma2Model,
         "Gemma3ForCausalLM": builder.Gemma3Model,
         "Gemma3ForConditionalGeneration": builder.Gemma3Model,
         "GraniteForCausalLM": builder.GraniteModel,

onnx_diagnostic/helpers/rt_helper.py

@@ -3,6 +3,7 @@ import numpy as np
 import onnx
 import torch
 from .helper import string_type, flatten_object
+from .torch_helper import to_numpy
 from .cache_helper import is_cache_dynamic_registered
 
 
@@ -56,7 +57,7 @@ def make_feeds(
         f"{string_type(torch.utils._pytree.tree_flatten(inputs)[0], with_shape=True)}"
     )
     if use_numpy:
-        flat = [t.detach().cpu().numpy() if isinstance(t, torch.Tensor) else t for t in flat]
+        flat = [to_numpy(t) if isinstance(t, torch.Tensor) else t for t in flat]
     names = (
         [i.name for i in proto.graph.input]
         if isinstance(proto, onnx.ModelProto)

onnx_diagnostic/helpers/torch_helper.py

@@ -5,7 +5,7 @@ import os
 import sys
 import warnings
 from collections.abc import Iterable
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 import numpy as np
 import onnx
 from onnx.external_data_helper import load_external_data_for_tensor, uses_external_data
@@ -283,9 +283,11 @@ def steal_forward(
     ],
     fprint: Callable = string_type,
     dump_file: Optional[str] = None,
+    dump_drop: Optional[Set[str]] = None,
     submodules: bool = False,
     verbose: int = 0,
     storage_limit: int = 2**27,
+    save_as_external_data: bool = True,
     **kwargs,
 ):
     """
@@ -303,6 +305,9 @@ def steal_forward(
     :param dump_file: dumps stolen inputs and outputs in an onnx model,
         they can be restored with :func:`create_input_tensors_from_onnx_model
         <onnx_diagnostic.helpers.mini_onnx_builder.create_input_tensors_from_onnx_model>`
+    :param dump_drop: to drop some inputs too big (only if dump_file is specified)
+    :param save_as_external_data: True by default, but maybe better to have everything
+        in a single file if possible
     :param submodules: if True and model is a module, the list extended with all the submodules
         the module contains
     :param verbose: verbosity
@@ -411,6 +416,15 @@ def steal_forward(
             if verbose:
                 size = torch_tensor_size(storage)
                 print(f"-- gather stored {len(storage)} objects, size={size // 2 ** 20} Mb")
+            if dump_drop:
+                for k, v in storage.items():
+                    if k[-1] == "I":
+                        _args, kwargs = v
+                        ii = set(kwargs) & dump_drop
+                        if ii:
+                            for i in ii:
+                                print("---", i)
+                                del kwargs[i]
             proto = create_onnx_model_from_input_tensors(storage)
             if verbose:
                 print("-- dumps stored objects")
@@ -420,7 +434,7 @@ def steal_forward(
             onnx.save(
                 proto,
                 dump_file,
-                save_as_external_data=True,
+                save_as_external_data=save_as_external_data,
                 all_tensors_to_one_file=True,
                 location=location,
             )
@@ -464,10 +478,10 @@ def is_torchdynamo_exporting() -> bool:
             return False
 
 
-def to_numpy(tensor: "torch.Tensor"):  # noqa: F821
+def to_numpy(tensor: "torch.Tensor") -> np.ndarray:  # noqa: F821
     """Converts a :class:`torch.Tensor` to :class:`numpy.ndarray`."""
     try:
-        return tensor.numpy()
+        return tensor.detach().cpu().numpy()
     except TypeError:
         # We try with ml_dtypes
         pass
@@ -476,7 +490,7 @@ def to_numpy(tensor: "torch.Tensor"):  # noqa: F821
 
     conv = {torch.bfloat16: ml_dtypes.bfloat16}
     assert tensor.dtype in conv, f"Unsupported type {tensor.dtype}, not in {conv}"
-    return tensor.to(torch.float32).numpy().astype(conv[tensor.dtype])
+    return tensor.detach().to(torch.float32).cpu().numpy().astype(conv[tensor.dtype])
 
 
 def replace_string_by_dynamic(dynamic_shapes: Any) -> Any:
@@ -765,7 +779,12 @@ def to_any(value: Any, to_value: Union[torch.dtype, torch.device, str]) -> Any:
 
 
 def torch_deepcopy(value: Any) -> Any:
-    """Makes a deepcopy."""
+    """
+    Makes a deep copy.
+
+    :param value: any value
+    :return: a deep copy
+    """
     if value is None:
         return None
     if isinstance(value, (int, float, str)):
@@ -794,9 +813,14 @@ def torch_deepcopy(value: Any) -> Any:
         from .cache_helper import CacheKeyValue
 
         ca = CacheKeyValue(value)
+        if len(ca.key_cache) == 0:
+            # Use of deepcopy.
+            import copy
+
+            return copy.deepcopy(value)
         return make_static_cache(
             torch_deepcopy(list(zip(ca.key_cache, ca.value_cache))),
-            max_cache_len=value.max_cache_len,
+            max_cache_len=max([value.max_cache_len, *[t.shape[2] for t in ca.key_cache]]),
         )
     if value.__class__.__name__ == "HybridCache":
         from .cache_helper import CacheKeyValue

onnx_diagnostic/reference/torch_evaluator.py

@@ -3,7 +3,7 @@ from typing import Dict, List, Optional, Sequence, Tuple, Union
 import numpy as np
 import onnx
 import torch
-from ..helpers.torch_helper import to_tensor
+from ..helpers.torch_helper import to_tensor, to_numpy
 from ..torch_onnx.runtime_info import first_used_last_used, RuntimeValue
 from .report_results_comparison import ReportResultComparison
 from . import torch_ops
@@ -578,7 +578,7 @@ class TorchOnnxEvaluator:
                 print(f"- clean {o}")
 
         if use_numpy:
-            return [None if a is None else a.detach().cpu().numpy() for a in fres]
+            return [None if a is None else to_numpy(a) for a in fres]
         return fres
 
     def run_with_values(

onnx_diagnostic/tasks/data/__init__.py

@@ -0,0 +1,13 @@
+import os
+
+
+def get_data(name: str):
+    """Returns data stored in this folder."""
+    filename = os.path.join(os.path.dirname(__file__), name)
+    assert os.path.exists(
+        filename
+    ), f"Unable to find a file with {name!r}, looked for {filename!r}"
+
+    from ...helpers.mini_onnx_builder import create_input_tensors_from_onnx_model
+
+    return create_input_tensors_from_onnx_model(filename)