onnx-diagnostic 0.6.2__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.2"
+__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,35 +186,57 @@ def get_parser_find() -> ArgumentParser:
         "--names",
         type=str,
         required=False,
-        help="names to look at comma separated values",
+        help="Names to look at comma separated values, if 'SHADOW', "
+        "search for shadowing names.",
     )
     parser.add_argument(
         "-v",
         "--verbose",
         default=0,
+        type=int,
         required=False,
         help="verbosity",
     )
+    parser.add_argument(
+        "--v2",
+        default=False,
+        action=BooleanOptionalAction,
+        help="Uses enumerate_results instead of onnx_find.",
+    )
     return parser
 
 
 def _cmd_find(argv: List[Any]):
-    from .helpers.onnx_helper import onnx_find
+    from .helpers.onnx_helper import onnx_find, enumerate_results, shadowing_names
 
     parser = get_parser_find()
     args = parser.parse_args(argv[1:])
-    onnx_find(args.input, verbose=args.verbose, watch=set(args.names.split(",")))
+    if args.names == "SHADOW":
+        onx = onnx.load(args.input, load_external_data=False)
+        s, ps = shadowing_names(onx)[:2]
+        print(f"shadowing names: {s}")
+        print(f"post-shadowing names: {ps}")
+    elif args.v2:
+        onx = onnx.load(args.input, load_external_data=False)
+        res = list(
+            enumerate_results(onx, name=set(args.names.split(",")), verbose=args.verbose)
+        )
+        if not args.verbose:
+            print("\n".join(map(str, res)))
+    else:
+        onnx_find(args.input, verbose=args.verbose, watch=set(args.names.split(",")))
 
 
 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(
@@ -227,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
@@ -270,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 {}
@@ -293,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")
@@ -319,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",
@@ -377,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(
@@ -402,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(
@@ -419,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()
@@ -471,6 +505,7 @@ def _cmd_validate(argv: List[Any]):
             runtime=args.runtime,
             repeat=args.repeat,
             warmup=args.warmup,
+            inputs2=args.inputs2,
         )
         print("")
         print("-- summary --")
@@ -481,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(
@@ -532,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
 
@@ -585,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

@@ -1,3 +1,29 @@
+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
@@ -30,3 +56,45 @@ def plot_legend(
     ax.grid(False)
     ax.set_axis_off()
     return ax
+
+
+def rotate_align(ax, angle=15, align="right"):
+    """Rotates x-label and align them to thr right. Returns ax."""
+    for label in ax.get_xticklabels():
+        label.set_rotation(angle)
+        label.set_horizontalalignment(align)
+    return ax
+
+
+def save_fig(ax, name: str):
+    """Applies ``tight_layout`` and saves the figures. Returns ax."""
+    import matplotlib.pyplot as plt
+
+    plt.tight_layout()
+    fig = ax.get_figure()
+    fig.savefig(name)
+    return ax
+
+
+def title(ax: "plt.axes", title: str) -> "plt.axes":  # noqa: F821
+    "Adds a title to axes and returns them."
+    ax.set_title(title)
+    return ax
+
+
+def plot_histogram(
+    tensor: np.ndarray,
+    ax: Optional["plt.axes"] = None,  # noqa: F821
+    bins: int = 30,
+    color: str = "orange",
+    alpha: float = 0.7,
+) -> "plt.axes":  # noqa: F821
+    "Computes the distribution for a tensor."
+    if ax is None:
+        import matplotlib.pyplot as plt
+
+        ax = plt.gca()
+        ax.cla()
+    ax.hist(tensor, bins=30, color="orange", alpha=0.7)
+    ax.set_yscale("log")
+    return ax

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/doc_helper.py

@@ -1,4 +1,5 @@
-from typing import Dict, Optional, Tuple
+import os
+from typing import Dict, List, Optional, Tuple
 import onnx
 import onnx.helper as oh
 import torch
@@ -6,6 +7,17 @@ from ..reference.torch_ops import OpRunKernel, OpRunTensor
 from .torch_helper import onnx_dtype_to_torch_dtype, torch_dtype_to_onnx_dtype
 from .ort_session import InferenceSessionForTorch
 
+_SAVED: List[str] = []
+_SAVE_OPTIMIZED_MODEL_ = int(os.environ.get("DUMP_ONNX", "0"))
+
+
+def _get_model_name(op_name: str, provider: str) -> Optional[str]:
+    if _SAVE_OPTIMIZED_MODEL_:
+        name = f"dump_doc_layer_norm_{provider}_{len(_SAVED)}.onnx"
+        _SAVED.append(name)
+        return name
+    return None
+
 
 class LayerNormalizationOrt(OpRunKernel):
     "LayerNormalization with onnxruntime"
@@ -13,14 +25,14 @@ class LayerNormalizationOrt(OpRunKernel):
     @classmethod
     def device_dependent(cls) -> bool:
         "Needs device."
-        return False
+        return True
 
     def __init__(
         self,
         node: onnx.NodeProto,
         version=None,
         device: Optional[torch.device] = None,
-        verbose=0,
+        verbose: int = 0,
     ):
         super().__init__(node, version, verbose=verbose)
         self.axis = self.get_attribute_int(node, "axis", -1)
@@ -70,7 +82,11 @@ class LayerNormalizationOrt(OpRunKernel):
         )
         provider = "CPUExecutionProvider" if self.is_cpu else "CUDAExecutionProvider"
         self._provider = provider
-        return InferenceSessionForTorch(layer_model, providers=[provider])
+        return InferenceSessionForTorch(
+            layer_model,
+            optimized_model_filepath=_get_model_name("layer_norm", provider),
+            providers=[provider],
+        )
 
     def run(self, x, scale, bias=None):
         itype = torch_dtype_to_onnx_dtype(x.dtype)
@@ -94,14 +110,14 @@ class MatMulOrt(OpRunKernel):
     @classmethod
     def device_dependent(cls) -> bool:
         "Needs device."
-        return False
+        return True
 
     def __init__(
         self,
         node: onnx.NodeProto,
         version=None,
         device: Optional[torch.device] = None,
-        verbose=0,
+        verbose: int = 0,
     ):
         super().__init__(node, version, verbose=verbose)
         self.device = device
@@ -127,7 +143,11 @@ class MatMulOrt(OpRunKernel):
         )
         provider = "CPUExecutionProvider" if self.is_cpu else "CUDAExecutionProvider"
         self._provider = provider
-        return InferenceSessionForTorch(model, providers=[provider])
+        return InferenceSessionForTorch(
+            model,
+            optimized_model_filepath=_get_model_name("matmul", provider),
+            providers=[provider],
+        )
 
     def run(self, a, b):
         itype = torch_dtype_to_onnx_dtype(a.dtype)