onnx-diagnostic 0.6.3__py3-none-any.whl → 0.7.0__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.6.3"
+__version__ = "0.7.0"
 __author__ = "Xavier Dupré"

onnx_diagnostic/_command_lines_parser.py

@@ -5,19 +5,18 @@ import re
 import sys
 import textwrap
 import onnx
-from typing import Any, List, Optional
+from typing import Any, Dict, List, Optional, Union
 from argparse import ArgumentParser, RawTextHelpFormatter, BooleanOptionalAction
-from textwrap import dedent
 
 
 def get_parser_lighten() -> ArgumentParser:
     parser = ArgumentParser(
         prog="lighten",
-        description=dedent(
+        description=textwrap.dedent(
+            """
+            Removes the weights from a heavy model, stores statistics to restore
+            random weights.
             """
-        Removes the weights from a heavy model, stores statistics to restore
-        random weights.
-        """
         ),
         epilog="This is mostly used to write unit tests without adding "
         "a big onnx file to the repository.",
@@ -70,11 +69,11 @@ def _cmd_lighten(argv: List[Any]):
 def get_parser_unlighten() -> ArgumentParser:
     parser = ArgumentParser(
         prog="unlighten",
-        description=dedent(
+        description=textwrap.dedent(
+            """
+            Restores random weights for a model reduces with command lighten,
+            the command expects to find a file nearby with extension '.stats'.
             """
-        Restores random weights for a model reduces with command lighten,
-        the command expects to find a file nearby with extension '.stats'.
-        """
         ),
         epilog="This is mostly used to write unit tests without adding "
         "a big onnx file to the repository.",
@@ -120,11 +119,7 @@ def _cmd_unlighten(argv: List[Any]):
 def get_parser_print() -> ArgumentParser:
     parser = ArgumentParser(
         prog="print",
-        description=dedent(
-            """
-        Prints the model on the standard output.
-        """
-        ),
+        description="Prints the model on the standard output.",
         epilog="To show a model.",
         formatter_class=RawTextHelpFormatter,
     )
@@ -171,11 +166,11 @@ def _cmd_print(argv: List[Any]):
 def get_parser_find() -> ArgumentParser:
     parser = ArgumentParser(
         prog="find",
-        description=dedent(
+        description=textwrap.dedent(
+            """
+            Look into a model and search for a set of names,
+            tells which node is consuming or producing it.
             """
-        Look into a model and search for a set of names,
-        tells which node is consuming or producing it.
-        """
         ),
         epilog="Enables Some quick validation.",
     )
@@ -191,8 +186,8 @@ def get_parser_find() -> ArgumentParser:
         "--names",
         type=str,
         required=False,
-        help="names to look at comma separated values, if 'SHADOW', "
-        "search for shadowing names",
+        help="Names to look at comma separated values, if 'SHADOW', "
+        "search for shadowing names.",
     )
     parser.add_argument(
         "-v",
@@ -206,7 +201,7 @@ def get_parser_find() -> ArgumentParser:
         "--v2",
         default=False,
         action=BooleanOptionalAction,
-        help="use enumerate_results instead of onnx_find",
+        help="Uses enumerate_results instead of onnx_find.",
     )
     return parser
 
@@ -235,12 +230,13 @@ def _cmd_find(argv: List[Any]):
 def get_parser_config() -> ArgumentParser:
     parser = ArgumentParser(
         prog="config",
-        description=dedent(
+        description=textwrap.dedent(
+            """
+            Prints out a configuration for a model id,
+            prints the associated task as well.
             """
-        Prints out a configuration for a model id,
-        prints the associated task as well.
-        """
         ),
+        formatter_class=RawTextHelpFormatter,
         epilog="",
     )
     parser.add_argument(
@@ -248,29 +244,29 @@ def get_parser_config() -> ArgumentParser:
         "--mid",
         type=str,
         required=True,
-        help="model id, usually <author>/<name>",
+        help="model id, usually `<author>/<name>`",
     )
     parser.add_argument(
         "-t",
         "--task",
         default=False,
         action=BooleanOptionalAction,
-        help="displays the task as well",
+        help="Displays the task as well.",
     )
     parser.add_argument(
         "-c",
         "--cached",
         default=True,
         action=BooleanOptionalAction,
-        help="uses cached configuration, only available for some of them, "
-        "mostly for unit test purposes",
+        help="Uses cached configuration, only available for some of them,\n"
+        "mostly for unit test purposes.",
     )
     parser.add_argument(
         "--mop",
         metavar="KEY=VALUE",
         nargs="*",
         help="Additional model options, use to change some parameters of the model, "
-        "example: --mop attn_implementation=eager",
+        "example:\n  --mop attn_implementation=sdpa or --mop attn_implementation=eager",
         action=_ParseDict,
     )
     return parser
@@ -291,6 +287,14 @@ def _cmd_config(argv: List[Any]):
         print(f"task: {task_from_id(args.mid)}")
 
 
+def _parse_json(value: str) -> Union[str, Dict[str, Any]]:
+    assert isinstance(value, str), f"value should be string but value={value!r}"
+    if value and value[0] == "{" and value[-1] == "}":
+        # a dictionary
+        return json.loads(value.replace("'", '"'))
+    return value
+
+
 class _ParseDict(argparse.Action):
     def __call__(self, parser, namespace, values, option_string=None):
         d = getattr(namespace, self.dest) or {}
@@ -314,22 +318,23 @@ class _ParseDict(argparse.Action):
                     continue
                 except (TypeError, ValueError):
                     pass
-                d[key] = value
+                d[key] = _parse_json(value)
 
         setattr(namespace, self.dest, d)
 
 
 def get_parser_validate() -> ArgumentParser:
     parser = ArgumentParser(
-        prog="test",
-        description=dedent(
+        prog="validate",
+        description=textwrap.dedent(
+            """
+            Prints out dummy inputs for a particular task or a model id.
+            If both mid and task are empty, the command line displays the list
+            of supported tasks.
             """
-        Prints out dummy inputs for a particular task or a model id.
-        If both mid and task are empty, the command line displays the list
-        of supported tasks.
-        """
         ),
         epilog="If the model id is specified, one untrained version of it is instantiated.",
+        formatter_class=RawTextHelpFormatter,
     )
     parser.add_argument("-m", "--mid", type=str, help="model id, usually <author>/<name>")
     parser.add_argument("-t", "--task", default=None, help="force the task to use")
@@ -340,55 +345,61 @@ def get_parser_validate() -> ArgumentParser:
         "--run",
         default=False,
         action=BooleanOptionalAction,
-        help="runs the model to check it runs",
+        help="Runs the model to check it runs.",
     )
     parser.add_argument(
         "-q",
         "--quiet",
         default=False,
         action=BooleanOptionalAction,
-        help="catches exception, report them in the summary",
+        help="Catches exception, reports them in the summary.",
     )
     parser.add_argument(
         "--patch",
         default=True,
         action=BooleanOptionalAction,
-        help="applies patches before exporting",
+        help="Applies patches before exporting.",
     )
     parser.add_argument(
         "--rewrite",
         default=True,
         action=BooleanOptionalAction,
-        help="applies rewrite before exporting",
+        help="Applies rewrite before exporting.",
     )
     parser.add_argument(
         "--stop-if-static",
         default=0,
         type=int,
-        help="raises an exception if a dynamic dimension becomes static",
+        help="Raises an exception if a dynamic dimension becomes static.",
     )
     parser.add_argument(
         "--trained",
         default=False,
         action=BooleanOptionalAction,
-        help="validate the trained model (requires downloading)",
+        help="Validates the trained model (requires downloading).",
+    )
+    parser.add_argument(
+        "--inputs2",
+        default=True,
+        action=BooleanOptionalAction,
+        help="Validates the model on a second set of inputs\n"
+        "to check the exported model supports dynamism.",
     )
     parser.add_argument(
         "--runtime",
         choices=["onnxruntime", "torch", "ref"],
         default="onnxruntime",
-        help="onnx runtime to use, onnxruntime by default",
+        help="onnx runtime to use, `onnxruntime` by default",
     )
     parser.add_argument(
         "-o",
         "--dump-folder",
-        help="if not empty, a folder is created to dumps statistics, "
-        "exported program, onnx...",
+        help="A folder is created to dumps statistics,\nexported program, onnx...",
     )
     parser.add_argument(
         "--drop",
-        help="drops the following inputs names, it should be a list "
-        "with comma separated values",
+        help="Drops the following inputs names, it should be a list\n"
+        "with comma separated values.",
     )
     parser.add_argument(
         "--opset",
@@ -398,24 +409,25 @@ def get_parser_validate() -> ArgumentParser:
     )
     parser.add_argument(
         "--subfolder",
-        help="subfolder where to find the model and the configuration",
+        help="Subfolder where to find the model and the configuration.",
     )
     parser.add_argument(
         "--ortfusiontype",
         required=False,
-        help="applies onnxruntime fusion, this parameter should contain the "
-        "model type or multiple values separated by `|`. `ALL` can be used "
-        "to run them all",
+        help="Applies onnxruntime fusion, this parameter should contain the\n"
+        "model type or multiple values separated by `|`. `ALL` can be used\n"
+        "to run them all.",
     )
     parser.add_argument("-v", "--verbose", default=0, type=int, help="verbosity")
-    parser.add_argument("--dtype", help="changes dtype if necessary")
-    parser.add_argument("--device", help="changes the device if necessary")
+    parser.add_argument("--dtype", help="Changes dtype if necessary.")
+    parser.add_argument("--device", help="Changes the device if necessary.")
     parser.add_argument(
         "--iop",
         metavar="KEY=VALUE",
         nargs="*",
-        help="Additional input options, use to change the default "
-        "inputs use to export, example: --iop cls_cache=SlidingWindowCache",
+        help="Additional input options, use to change the default"
+        "inputs use to export, example:\n  --iop cls_cache=SlidingWindowCache"
+        "\n  --iop cls_cache=StaticCache",
         action=_ParseDict,
     )
     parser.add_argument(
@@ -423,7 +435,8 @@ def get_parser_validate() -> ArgumentParser:
         metavar="KEY=VALUE",
         nargs="*",
         help="Additional model options, use to change some parameters of the model, "
-        "example: --mop attn_implementation=eager",
+        "example:\n  --mop attn_implementation=sdpa --mop attn_implementation=eager\n  "
+        "--mop \"rope_scaling={'rope_type': 'dynamic', 'factor': 10.0}\"",
         action=_ParseDict,
     )
     parser.add_argument(
@@ -440,7 +453,7 @@ def get_parser_validate() -> ArgumentParser:
 
 def _cmd_validate(argv: List[Any]):
     from .helpers import string_type
-    from .torch_models.test_helper import get_inputs_for_task, validate_model
+    from .torch_models.validate import get_inputs_for_task, validate_model
     from .tasks import supported_tasks
 
     parser = get_parser_validate()
@@ -492,6 +505,7 @@ def _cmd_validate(argv: List[Any]):
             runtime=args.runtime,
             repeat=args.repeat,
             warmup=args.warmup,
+            inputs2=args.inputs2,
         )
         print("")
         print("-- summary --")
@@ -502,11 +516,7 @@ def _cmd_validate(argv: List[Any]):
 def get_parser_stats() -> ArgumentParser:
     parser = ArgumentParser(
         prog="stats",
-        description=dedent(
-            """
-        Prints out statistics on an ONNX model.
-        """
-        ),
+        description="Prints out statistics on an ONNX model.",
         epilog="",
     )
     parser.add_argument(
@@ -553,8 +563,8 @@ def get_parser_stats() -> ArgumentParser:
         required=False,
         default="",
         type=str,
-        help="keeps only tensors whose name verifies "
-        "this regular expression, empty = no filter",
+        help="Keeps only tensors whose name verifies "
+        "this regular expression, empty = no filter.",
     )
     return parser
 
@@ -606,17 +616,17 @@ def get_main_parser() -> ArgumentParser:
         formatter_class=RawTextHelpFormatter,
         epilog=textwrap.dedent(
             """
-        Type 'python -m onnx_diagnostic <cmd> --help'
-        to get help for a specific command.
-
-        config     - prints a configuration for a model id
-        find       - find node consuming or producing a result
-        lighten    - makes an onnx model lighter by removing the weights,
-        unlighten  - restores an onnx model produces by the previous experiment
-        print      - prints the model on standard output
-        validate   - validate a model
-        stats      - produces statistics on a model
-        """
+            Type 'python -m onnx_diagnostic <cmd> --help'
+            to get help for a specific command.
+
+            config     - prints a configuration for a model id
+            find       - find node consuming or producing a result
+            lighten    - makes an onnx model lighter by removing the weights,
+            unlighten  - restores an onnx model produces by the previous experiment
+            print      - prints the model on standard output
+            validate   - validate a model
+            stats      - produces statistics on a model
+            """
         ),
     )
     parser.add_argument(

onnx_diagnostic/doc.py

@@ -2,6 +2,28 @@ from typing import Optional
 import numpy as np
 
 
+def get_latest_pypi_version(package_name="onnx-diagnostic") -> str:
+    """Returns the latest published version."""
+
+    import requests
+
+    url = f"https://pypi.org/pypi/{package_name}/json"
+    response = requests.get(url)
+
+    assert response.status_code == 200, f"Unable to retrieve the version response={response}"
+    data = response.json()
+    version = data["info"]["version"]
+    return version
+
+
+def update_version_package(version: str, package_name="onnx-diagnostic") -> str:
+    "Adds dev if the major version is different from the latest published one."
+    released = get_latest_pypi_version(package_name)
+    shorten_r = ".".join(released.split(".")[:2])
+    shorten_v = ".".join(version.split(".")[:2])
+    return version if shorten_r == shorten_v else f"{shorten_v}.dev"
+
+
 def reset_torch_transformers(gallery_conf, fname):
     "Resets torch dynamo for :epkg:`sphinx-gallery`."
     import matplotlib.pyplot as plt

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

@@ -141,6 +141,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],