onnx-diagnostic 0.8.0__py3-none-any.whl

onnx_diagnostic/torch_export_patches/patch_expressions.py

@@ -0,0 +1,108 @@
+from typing import Callable, Set
+import torch
+from ..helpers.torch_helper import is_torchdynamo_exporting
+
+
+def make_undefined_dimension(i: int) -> torch.SymInt:
+    """
+    Uses for a custom op when a new dimension must be introduced to bypass
+    some verification. The following function creates a dummy output
+    with a dimension based on the content.
+
+    .. code-block:: python
+
+        def symbolic_shape(x, y):
+            return torch.empty(
+                x.shape[0],
+                make_undefined_dimension(min(x.shape[1], y[0])),
+            )
+    """
+    try:
+        ti = int(i)
+    except:  # noqa: E722
+        ti = 10
+    t = torch.ones((ti * 2,))
+    t[:ti] = 0
+    res = torch.nonzero(t).shape[0]
+    return res
+
+
+def _patched_float_arange(
+    start: torch.Tensor, end: torch.Tensor, step: torch.Tensor
+) -> torch.Tensor:
+    """Float arange."""
+    return torch.arange(
+        float(start.item()),
+        float(end.item()),
+        float(step.item()),
+        dtype=start.dtype,
+        device=start.device,
+    )
+
+
+def _patched_float_arange_shape(start, end, step):
+    # Fails because:
+    # Did you accidentally call new_dynamic_size() or item()
+    # more times than you needed to in your fake implementation?
+    # try:
+    #     n = math.ceil(((end - start) / step).item())
+    # except:  # noqa: E722
+    #     n = 10
+    n = 10
+    return torch.empty((make_undefined_dimension(n),), dtype=start.dtype, device=start.device)
+
+
+def _iterate_patched_expressions():
+    glo = globals().copy()
+    for k, _v in glo.items():
+        if k.startswith("_patched_") and not k.endswith("_shape"):
+            name = k
+            yield k[len("_patched_") :], glo[name], glo[f"{name}_shape"]
+
+
+_registered: Set[str] = set()
+
+
+def _register_patched_expression(
+    fct: Callable, fct_shape: Callable, namespace: str, fname: str
+):
+    schema_str = torch.library.infer_schema(fct, mutates_args=())
+    custom_def = torch.library.CustomOpDef(namespace, fname, schema_str, fct)
+    custom_def.register_kernel("cpu")(fct)
+    custom_def._abstract_fn = fct_shape
+
+
+def register_patched_expressions(namespace: str = "patched"):
+    """
+    Registers as custom ops known expressions failing due to dynamic shapes.
+
+    .. runpython::
+        :showcode:
+
+        import pprint
+        from onnx_diagnostic.torch_export_patches.patch_expressions import (
+            _iterate_patched_expressions,
+        )
+
+        pprint.pprint([name for name, _f, _fsh in _iterate_patched_expressions()])
+    """
+    for name, f, fsh in _iterate_patched_expressions():
+        if name not in _registered:
+            _register_patched_expression(f, fsh, namespace, name)
+            _registered.add(name)
+
+
+def patched_selector(fct: Callable, patched_fct: Callable) -> Callable:
+    """
+    Returns **fct** if the model is being executed or
+    **patched_fct** if it is being exported.
+    """
+    return patched_fct if is_torchdynamo_exporting() else fct
+
+
+def patched_float_arange(start, end, step):
+    """Patched arange when start, end, step are floats."""
+    if is_torchdynamo_exporting():
+        return torch.ops.patched.float_arange(start, end, step)
+    else:
+        return torch.arange(start, end, step)

onnx_diagnostic/torch_export_patches/patch_inputs.py

@@ -0,0 +1,211 @@
+import inspect
+from typing import Any, Dict, Optional, Tuple
+import torch
+import transformers
+from ..helpers import string_type
+
+
+def _process_cache(k: str, v):
+    assert k != "position_ids" or isinstance(
+        k, torch.Tensor
+    ), f"Unexpected type for parameter {k!r} {string_type(v, with_shape=True)}"
+    if (
+        isinstance(v, list)
+        and all(isinstance(i, tuple) for i in v)
+        and set(len(t) for t in v) == {2}
+    ):
+        # A dynamicCache
+        from ..helpers.cache_helper import make_dynamic_cache
+
+        cache = make_dynamic_cache(v)
+        return cache
+    if isinstance(v, torch.Tensor):
+        return v
+    raise NotImplementedError(
+        f"Unable to process parameter {k!r} with v={string_type(v,with_shape=True)}"
+    )
+
+
+def _make_shape(subset: Dict, cls: type, value: Any) -> Any:
+    if cls is transformers.cache_utils.DynamicCache:
+        assert subset, "DynamicCache cannot be empty"
+        values = set(map(str, subset.values()))
+        assert len(values) == 1, (
+            f"Inconsistencies in subset={subset}, found={values}, "
+            f"it cannot be a {cls}, value={string_type(value)}"
+        )
+        cache_length = len(value.layers if hasattr(value, "layers") else value.key_cache)
+        for v in subset.values():
+            axes = v
+            break
+        new_shape = [axes for i in range(cache_length * 2)]
+        return new_shape
+    if value.__class__ in torch.utils._pytree.SUPPORTED_NODES:
+        raise NotImplementedError(
+            f"_make_shape not implemented for registered class={cls}, "
+            f"subset={subset}, value={string_type(value)}"
+        )
+    raise NotImplementedError(
+        f"_make_shape not implemented for cls={cls}, "
+        f"subset={subset}, value={string_type(value)}"
+    )
+
+
+def convert_dynamic_axes_into_dynamic_shapes(
+    model: torch.nn.Module,
+    args: Optional[Tuple[Any, ...]] = None,
+    kwargs: Optional[Dict[str, Any]] = None,
+    dynamic_axes: Optional[Dict[str, Dict[int, str]]] = None,
+    prefix_mapping: Optional[Dict[str, str]] = None,
+    verbose: int = 0,
+) -> Tuple[Tuple[Any, ...], Dict[str, Any], Dict[str, Any]]:
+    """
+    Converts the input from an export to something :func:`torch.export.export` can handle.
+
+    :param model: model to convert (used to extract the signature)
+    :param args: positional arguments
+    :param kwargs: named arguments
+    :param dynamic_axes: dynamic axes
+    :param prefix_mapping: prefix mapping
+    :param verbose: verbosity
+    :return: (args, kwargs, dynamic shapes)
+    """
+    from ..helpers.cache_helper import CacheKeyValue
+
+    new_kwargs = {}
+    if args:
+        assert hasattr(model, "forward"), f"Missing method 'forward' for {model!r}"
+        plus = 0 if isinstance(model, torch.nn.Module) else 1
+        print(
+            f"[convert_dynamic_axes_into_dynamic_shapes] "
+            f"mapping args to kwargs for model="
+            f"{model if plus else model.__class__.__name__}"
+        )
+        pars = inspect.signature(model.forward).parameters
+        assert len(pars) >= len(
+            args
+        ), f"Length mismatch, len(args)={len(args)}, pars={list(pars)}"
+
+        for i, p in enumerate(pars):
+            if i < plus:
+                continue
+            if i - plus >= len(args):
+                break
+            if verbose:
+                print(
+                    f"[convert_dynamic_axes_into_dynamic_shapes] mapping args[{i-plus}] "
+                    f"to {p!r} ({string_type(args[i-plus])})"
+                )
+            new_kwargs[p] = args[i - plus]
+
+    if kwargs:
+        for k, v in kwargs.items():
+            assert k not in new_kwargs, f"Argument {k!r} from kwargs already present in args."
+            new_kwargs[k] = v
+
+    # process
+    updated_kwargs = {}
+    changes = {}
+    for k, v in new_kwargs.items():
+        if isinstance(v, torch.Tensor):
+            updated_kwargs[k] = v
+            continue
+        if isinstance(v, list):
+            # cache?
+            updated_kwargs[k] = _process_cache(k, v)
+            if type(updated_kwargs[k]) is not type(v):
+                # A cache was introduced.
+                if verbose:
+                    print(
+                        f"[convert_dynamic_axes_into_dynamic_shapes] parameter "
+                        f"{k!r} was changed into {type(updated_kwargs[k])}"
+                    )
+                changes[k] = type(updated_kwargs[k])
+                continue
+        if isinstance(v, transformers.cache_utils.DynamicCache):
+            ca = CacheKeyValue(v)
+            updated_kwargs[k] = [ca.key_cache, ca.value_cache]
+            changes[k] = type(v)
+            continue
+        raise NotImplementedError(
+            f"Unexpected type {type(v)} for parameter {k!r} "
+            f"({string_type(v, with_shape=True)})"
+        )
+
+    # process dynamic axes
+    if changes:
+        dynamic_shapes = {}
+        done = set()
+        for k, v in dynamic_axes.items():
+            if k not in changes and k in updated_kwargs and isinstance(v, dict):
+                dynamic_shapes[k] = v
+                continue
+            if (
+                k in updated_kwargs
+                and k in changes
+                and changes[k] == transformers.cache_utils.DynamicCache
+            ):
+                dynamic_shapes[k] = v
+                continue
+            if "." in k:
+                # something like present.0.key
+                prefix = k.split(".")[0]
+                if prefix in done:
+                    continue
+                args_prefix = (
+                    prefix_mapping[prefix]
+                    if prefix_mapping and prefix in prefix_mapping
+                    else prefix
+                )
+                if args_prefix in updated_kwargs and args_prefix in changes:
+                    # A cache.
+                    cls = changes[args_prefix]
+                    dynamic_shapes[args_prefix] = _make_shape(
+                        {
+                            _: __
+                            for _, __ in dynamic_axes.items()
+                            if _.startswith(f"{prefix}.")
+                        },
+                        cls,
+                        updated_kwargs[args_prefix],
+                    )
+                    done.add(prefix)
+                    continue
+            if k not in updated_kwargs:
+                # dynamic axes not in the given inputs, should be raise an exception?
+                if verbose:
+                    print(
+                        f"[convert_dynamic_axes_into_dynamic_shapes] dropping axes "
+                        f"{k!r}-{v!r}, not found in {set(updated_kwargs)}"
+                    )
+                continue
+            raise NotImplementedError(
+                f"Unable to process dynamic axes {k!r}, axes={v}, "
+                f"value={string_type(updated_kwargs[k], with_shape=True)}, "
+                f"dynamic axes={dynamic_axes}, "
+                f"updated_kwargs={string_type(updated_kwargs, with_shape=True)}"
+            )
+
+    return (), updated_kwargs, dynamic_shapes
+
+
+def use_dyn_not_str(dynamic_shapes: Any, default_value=None) -> Any:
+    """
+    Some functions returns dynamic shapes as string.
+    This functions replaces them with ``torch.export.Dim.DYNAMIC``.
+    ``default_value=torch.export.Dim.AUTO`` changes the default value.
+    """
+    if isinstance(dynamic_shapes, list):
+        return [use_dyn_not_str(a, default_value=default_value) for a in dynamic_shapes]
+    if isinstance(dynamic_shapes, tuple):
+        return tuple(use_dyn_not_str(a, default_value=default_value) for a in dynamic_shapes)
+    if isinstance(dynamic_shapes, dict):
+        return {
+            k: use_dyn_not_str(v, default_value=default_value)
+            for k, v in dynamic_shapes.items()
+        }
+    if isinstance(dynamic_shapes, set):
+        return {use_dyn_not_str(a, default_value=default_value) for a in dynamic_shapes}
+    if isinstance(dynamic_shapes, str):
+        return torch.export.Dim.DYNAMIC if default_value is None else default_value
+    return dynamic_shapes