onnx-diagnostic 0.6.3__py3-none-any.whl → 0.7.1__py3-none-any.whl

onnx_diagnostic/export/dynamic_shapes.py

@@ -630,9 +630,12 @@ class ModelInputs:
         method_name: str = "forward",
         name: str = "main",
     ):
-        assert isinstance(model, torch.nn.Module) or inspect.ismodule(
-            model
-        ), f"unexpected type for model={type(model)}, it must be a torch.nn.Module"
+        assert (
+            model is None or isinstance(model, torch.nn.Module) or inspect.ismodule(model)
+        ), (
+            f"unexpected type for model={type(model)}, "
+            f"it must be a torch.nn.Module or None"
+        )
         assert name, (
             f"name={name!r} cannot be empty this string is used to "
             f"display meaningful error messages"
@@ -641,26 +644,42 @@ class ModelInputs:
         self.model = model
         self.level = level
         self.method_name = method_name
-        self.forward = getattr(model, method_name)
-        self.signature = inspect.signature(self.forward)
+        self.forward = getattr(model, method_name) if model is not None else None
+        self.signature = inspect.signature(self.forward) if self.forward else None
 
         # information about the signature
-        self.forward_parameter_names = set(
-            p.name
-            for p in self.signature.parameters.values()
-            if p.kind not in {p.VAR_POSITIONAL, p.VAR_KEYWORD}
+        self.forward_parameter_names = (
+            set(
+                p.name
+                for p in self.signature.parameters.values()
+                if p.kind not in {p.VAR_POSITIONAL, p.VAR_KEYWORD}
+            )
+            if self.signature
+            else None
+        )
+        self.forward_ordered_parameter_names = (
+            list(self.signature.parameters) if self.signature else None
+        )
+        self.forward_positioned_parameter_names = (
+            [
+                p.name
+                for p in self.signature.parameters.values()
+                if p.kind in (p.VAR_POSITIONAL, p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD)
+            ]
+            if self.signature
+            else None
+        )
+        names = (
+            [p.name for p in self.signature.parameters.values() if p.kind == p.VAR_POSITIONAL]
+            if self.signature
+            else None
         )
-        self.forward_ordered_parameter_names = list(self.signature.parameters)
-        self.forward_positioned_parameter_names = [
-            p.name
-            for p in self.signature.parameters.values()
-            if p.kind in (p.VAR_POSITIONAL, p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD)
-        ]
-        names = [
-            p.name for p in self.signature.parameters.values() if p.kind == p.VAR_POSITIONAL
-        ]
         self.forward_args = names[0] if names else None
-        names = [p.name for p in self.signature.parameters.values() if p.kind == p.VAR_KEYWORD]
+        names = (
+            [p.name for p in self.signature.parameters.values() if p.kind == p.VAR_KEYWORD]
+            if self.signature
+            else None
+        )
         self.forward_kwargs = names[0] if names else None
         self.forward_custom_op_schema = None
         self.forward_need_serialization = False
@@ -711,6 +730,7 @@ class ModelInputs:
     @property
     def true_model_name(self) -> str:
         "Returns class name or module name."
+        assert self.model is not None, "model was None when the class was initialized."
         return (
             self.model.__class__.__name__
             if isinstance(self.model, torch.nn.Module)
@@ -942,7 +962,7 @@ class ModelInputs:
                 )
             )
         names = s2.pop()
-        for name in names:
+        for i, name in enumerate(names):
             assert name not in {"_diag", "verbose"}, (
                 f"{self.full_name}: unexpected parameter {name!r}, names={names}"
                 f"\ninputs[0]={string_type(self.inputs[0], with_shape=True)}"
@@ -968,6 +988,14 @@ class ModelInputs:
         with the corresponding dynamic shapes.
         *kwargs*, *dynamic_shapes* are modified inplace.
         """
+        assert (
+            self.signature is not None
+            and self.forward_parameter_names is not None
+            and self.forward_ordered_parameter_names is not None
+        ), (
+            "model was None when the class was initialized, "
+            "cannot move args to kwargs without the signature."
+        )
         sig = self.signature
         arg_dyn, kw_dyn = dynamic_shapes
         for i, p in enumerate(sig.parameters):

onnx_diagnostic/export/shape_helper.py

@@ -0,0 +1,126 @@
+from typing import Any, Dict, List, Set, Tuple, Union
+from ..helpers.cache_helper import flatten_unflatten_for_dynamic_shapes
+from .dynamic_shapes import ModelInputs
+
+
+def all_dynamic_shape_from_inputs(inputs: Any, dim_prefix: Any = "d") -> Any:
+    """
+    Returns the dynamic shapes for the given inputs.
+    All dimensions are considered as dynamic.
+    ``dim_prefix`` can be a string (the function uses it as a prefix),
+    or ``torch.export.Dim.AUTO`` or ``torch.export.Dim.DYNAMIC``.
+
+    .. runpython::
+        :showcode:
+
+        import pprint
+        import torch
+        from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
+        from onnx_diagnostic.export.shape_helper import all_dynamic_shape_from_inputs
+
+        bsize, nheads, slen, dim = 2, 1, 30, 96
+        inputs = dict(
+            input_ids=torch.randint(15, size=(2, 3), dtype=torch.int64),
+            attention_mask=torch.randint(1, size=(2, 33), dtype=torch.int64),
+            position_ids=torch.arange(3, dtype=torch.int64),
+            past_key_values=make_dynamic_cache(
+                [(torch.randn(bsize, nheads, slen, dim),
+                  torch.randn(bsize, nheads, slen, dim))]
+            ),
+        )
+        ds = all_dynamic_shape_from_inputs(inputs)
+        pprint.pprint(ds)
+    """
+    if isinstance(dim_prefix, str):
+        prefixes: Set[str] = set()
+
+        def tensor_to_shape(tensor):
+            n = len(prefixes)
+            p = f"{dim_prefix}_{n}"
+            prefixes.add(p)
+            return {i: f"{p}_{i}" for i in range(tensor.ndim)}
+
+    else:
+
+        def tensor_to_shape(tensor):
+            return {i: dim_prefix for i in range(tensor.ndim)}  # noqa: C420
+
+    return flatten_unflatten_for_dynamic_shapes(
+        inputs, change_function=tensor_to_shape, use_dict=True
+    )
+
+
+def guess_dynamic_shapes_from_inputs(
+    inputs: List[Any], auto: Union[bool, str] = False
+) -> Tuple[Tuple[Any, ...], Dict[str, Any]]:
+    """
+    Guesses which dimension is dimension from a set of inputs.
+    Every dimension having different values over multiple sets
+    of inputs. Every dimension not changing remains static.
+
+    :param inputs: a list of input sets
+    :param auto: True for ``torch.export.Dim.AUTO``,
+        False for ``torch.export.Dim.DYNAMIC``,
+        a string to get a unique string for every dynamic dimension
+    :return: args and kwargs
+
+    .. runpython::
+        :showcode:
+
+        import pprint
+        import torch
+        from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
+        from onnx_diagnostic.export.shape_helper import guess_dynamic_shapes_from_inputs
+
+        bsize, nheads, slen, dim = 2, 1, 30, 96
+        inputs1 = dict(
+            input_ids=torch.randint(15, size=(2, 3), dtype=torch.int64),
+            attention_mask=torch.randint(1, size=(2, 33), dtype=torch.int64),
+            position_ids=torch.arange(3, dtype=torch.int64),
+            past_key_values=make_dynamic_cache(
+                [
+                    (
+                        torch.randn(bsize, nheads, slen, dim),
+                        torch.randn(bsize, nheads, slen, dim),
+                    ),
+                ]
+            ),
+        )
+        bsize, nheads, slen, dim = 3, 1, 33, 96
+        inputs2 = dict(
+            input_ids=torch.randint(15, size=(3, 4), dtype=torch.int64),
+            attention_mask=torch.randint(1, size=(3, 34), dtype=torch.int64),
+            position_ids=torch.arange(4, dtype=torch.int64),
+            past_key_values=make_dynamic_cache(
+                [
+                    (
+                        torch.randn(bsize, nheads, slen, dim),
+                        torch.randn(bsize, nheads, slen, dim),
+                    ),
+                ]
+            ),
+        )
+        ds = guess_dynamic_shapes_from_inputs([inputs1, inputs2], auto="d")
+        pprint.pprint(ds)
+
+    This function returns something equivalent to function
+    :class:`torch.export.dynamic_shapes.AdditionalInputs` but this
+    one needs a model.
+
+    .. runpython::
+        :showcode:
+
+        import pprint
+        import torch
+        from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
+        from onnx_diagnostic.export.shape_helper import guess_dynamic_shapes_from_inputs
+        from onnx_diagnostic.torch_models.hghub import get_untrained_model_with_inputs
+
+        data = get_untrained_model_with_inputs("arnir0/Tiny-LLM", add_second_input=True)
+        ds = torch.export.dynamic_shapes.AdditionalInputs()
+        ds.add((), data["inputs"])
+        ds.add((), data["inputs2"])
+        pprint.pprint(ds.dynamic_shapes(data["model"], (), data["inputs"]))
+    """
+    mi = ModelInputs(None, inputs)
+    return mi.guess_dynamic_shapes(auto=auto)

onnx_diagnostic/ext_test_case.py

@@ -1014,7 +1014,7 @@ class ExtTestCase(unittest.TestCase):
                     msg_ = "\n".join(excs)
                     msg = f"{msg}\n{msg_}" if msg else msg_
                     raise AssertionError(f"Found {len(excs)} discrepancies\n{msg}")
-        elif expected.__class__.__name__ == "DynamicCache":
+        elif expected.__class__.__name__ in ("DynamicCache", "StaticCache"):
             atts = {"key_cache", "value_cache"}
             self.assertEqualArrayAny(
                 {k: expected.__dict__.get(k, None) for k in atts},

onnx_diagnostic/helpers/cache_helper.py

@@ -1,11 +1,15 @@
-from typing import Any, List, Tuple
+from typing import Any, Callable, List, Optional, Tuple
 import packaging.version as pv
 import torch
 import transformers
 import transformers.cache_utils
 
 
-def flatten_unflatten_for_dynamic_shapes(obj: Any, use_dict: bool = False) -> Any:
+def flatten_unflatten_for_dynamic_shapes(
+    obj: Any,
+    use_dict: bool = False,
+    change_function: Optional[Callable[[torch.Tensor], Any]] = None,
+) -> Any:
     """
     Returns the object in a different structure similar to what
     the definition of the dynamic shapes should use.
@@ -15,11 +19,13 @@ def flatten_unflatten_for_dynamic_shapes(obj: Any, use_dict: bool = False) -> An
         :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.
+        for the strict and non strict mode. It also preserves tuple.
+    :param change_function: to modifies the tensor in the structure itself,
+        like replace them by a shape
     :return: the serialized object
     """
     if isinstance(obj, torch.Tensor):
-        return obj
+        return change_function(obj) if change_function else obj
     flat, spec = torch.utils._pytree.tree_flatten(obj)
     start = 0
     end = 0
@@ -27,12 +33,17 @@ def flatten_unflatten_for_dynamic_shapes(obj: Any, use_dict: bool = False) -> An
     for subspec in spec.children_specs:
         end += subspec.num_leaves
         value = subspec.unflatten(flat[start:end])
-        value = flatten_unflatten_for_dynamic_shapes(value, use_dict=use_dict)
+        value = flatten_unflatten_for_dynamic_shapes(
+            value, use_dict=use_dict, change_function=change_function
+        )
         subtrees.append(value)
         start = end
-    if use_dict and (spec.type is dict or spec.context):
-        # This a dictionary.
-        return dict(zip(spec.context, subtrees))
+    if use_dict:
+        if spec.type is dict or spec.context:
+            # This a dictionary.
+            return dict(zip(spec.context, subtrees))
+        if spec.type is tuple:
+            return tuple(subtrees)
     # This is a list.
     return subtrees
 
@@ -141,6 +152,65 @@ else:
         return cache
 
 
+def make_static_cache(
+    key_value_pairs: List[Tuple[torch.Tensor, torch.Tensor]],
+) -> transformers.cache_utils.DynamicCache:
+    """
+    Creates an instance of :class:`transformers.cache_utils.StaticCache`.
+    :param key_value_pairs: list of pairs of (key, values)
+    :return: :class:`transformers.cache_utils.StaticCache`
+
+    Example:
+
+    .. runpython::
+        :showcode:
+
+        import torch
+        from onnx_diagnostic.helpers import string_type
+        from onnx_diagnostic.helpers.cache_helper import make_static_cache
+
+        n_layers = 2
+        bsize, nheads, slen, dim = 2, 4, 3, 7
+
+        past_key_values = make_static_cache(
+            [
+                (
+                    torch.randn(bsize, nheads, slen, dim),
+                    torch.randn(bsize, nheads, slen, dim),
+                )
+                for i in range(n_layers)
+            ]
+        )
+        print(string_type(past_key_values, with_shape=True))
+    """
+
+    class _config:
+        def __init__(self):
+            self.head_dim = key_value_pairs[0][0].shape[-1]
+            self.num_attention_heads = key_value_pairs[0][0].shape[1]
+            self.num_hidden_layers = len(key_value_pairs)
+
+    cache = transformers.cache_utils.StaticCache(
+        _config(),
+        max_batch_size=key_value_pairs[0][0].shape[0],
+        device=key_value_pairs[0][0].device,
+        dtype=key_value_pairs[0][0].dtype,
+        max_cache_len=key_value_pairs[0][0].shape[2],
+    )
+    for i in range(len(key_value_pairs)):
+        assert cache.key_cache[i].shape == key_value_pairs[i][0].shape, (
+            f"Shape mismatch, expected {cache.key_cache[i].shape}, "
+            f"got {key_value_pairs[i][0].shape}"
+        )
+        cache.key_cache[i][:, :, :, :] = key_value_pairs[i][0]
+        assert cache.value_cache[i].shape == key_value_pairs[i][1].shape, (
+            f"Shape mismatch, expected {cache.value_cache[i].shape}, "
+            f"got {key_value_pairs[i][1].shape}"
+        )
+        cache.value_cache[i][:, :, :, :] = key_value_pairs[i][1]
+    return cache
+
+
 def make_encoder_decoder_cache(
     self_attention_cache: transformers.cache_utils.DynamicCache,
     cross_attention_cache: transformers.cache_utils.DynamicCache,

onnx_diagnostic/helpers/config_helper.py

@@ -34,10 +34,14 @@ def update_config(config: Any, mkwargs: Dict[str, Any]):
                 config._attn_implementation_autoset = False
             continue
         if isinstance(v, dict):
-            assert hasattr(
-                config, k
-            ), f"missing attribute {k!r} in config={config}, cannot update it with {v}"
-            update_config(getattr(config, k), v)
+            if not hasattr(config, k) or getattr(config, k) is None:
+                setattr(config, k, v)
+                continue
+            existing = getattr(config, k)
+            if type(existing) is dict:
+                existing.update(v)
+            else:
+                update_config(getattr(config, k), v)
             continue
         setattr(config, k, v)
 

onnx_diagnostic/helpers/helper.py

@@ -558,7 +558,7 @@ def string_type(
             print(f"[string_type] CACHE1:{type(obj)}")
         return f"MambaCache(conv_states={c}, ssm_states={d})"
 
-    if obj.__class__.__name__ in ("DynamicCache", "SlidingWindowCache"):
+    if obj.__class__.__name__ in {"DynamicCache", "SlidingWindowCache", "StaticCache"}:
         kc = string_type(
             obj.key_cache,
             with_shape=with_shape,
@@ -857,7 +857,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__ == "DynamicCache":
+    if x.__class__.__name__ in {"DynamicCache", "StaticCache"}:
         res = flatten_object(x.key_cache) + flatten_object(x.value_cache)
         return tuple(res)
     if x.__class__.__name__ == "EncoderDecoderCache":
@@ -1424,10 +1424,37 @@ def max_diff(
             f"level={level}"
         )
 
+    if expected.__class__.__name__ == "StaticCache":
+        if got.__class__.__name__ == "StaticCache":
+            if verbose >= 6:
+                print(f"[max_diff] StaticCache: {string_type(expected)} ? {string_type(got)}")
+            return max_diff(
+                [expected.key_cache, expected.value_cache],
+                [got.key_cache, got.value_cache],
+                verbose=verbose,
+                hist=hist,
+            )
+        if isinstance(got, tuple) and len(got) == 2:
+            return max_diff(
+                [expected.key_cache, expected.value_cache],
+                [got[0], got[1]],
+                debug_info=_debug(expected.__class__.__name__),
+                **_dkws,
+            )
+        raise AssertionError(
+            f"StaticCache not fully implemented with classes "
+            f"{expected.__class__.__name__!r} and {got.__class__.__name__!r}, "
+            f"and expected={string_type(expected)}, got={string_type(got)},\n"
+            f"level={level}"
+        )
+
     if expected.__class__.__name__ == "SlidingWindowCache":
         if got.__class__.__name__ == "SlidingWindowCache":
             if verbose >= 6:
-                print(f"[max_diff] DynamicCache: {string_type(expected)} ? {string_type(got)}")
+                print(
+                    f"[max_diff] SlidingWindowCache: "
+                    f"{string_type(expected)} ? {string_type(got)}"
+                )
             return max_diff(
                 [expected.key_cache, expected.value_cache],
                 [got.key_cache, got.value_cache],