onnx-diagnostic 0.7.16__py3-none-any.whl → 0.8.1__py3-none-any.whl

onnx_diagnostic/helpers/fake_tensor_helper.py

@@ -1,82 +1,236 @@
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, Optional, Set, Tuple
 
 
-_UNIQUE = set()
+class FakeTensorContext:
+    """Stores information used to reused same dimension for the same dimension names."""
 
+    def __init__(self, fake_mode: Optional["FakeTensorMode"] = None):  # noqa: F821
+        if fake_mode is None:
+            from torch.fx.experimental.symbolic_shapes import ShapeEnv
+            from torch._subclasses.fake_tensor import FakeTensorMode
 
-def _unique():
-    i = 129 + 1
-    while i in _UNIQUE:
-        i += 1
-    _UNIQUE.add(i)
-    return i
+            shape_env = ShapeEnv()
+            self.fake_mode = FakeTensorMode(shape_env=shape_env)
+        else:
+            self.fake_mode = fake_mode
+        self._candidates = self._first_primes()
+        self._unique_: Set[str] = set()
+        self._mapping_int: Dict[int, str] = {}
+        self._mapping_str: Dict[str, int] = {}
 
+    @classmethod
+    def _first_primes(cls, n=1000):
+        sieve = [True] * (n + 1)
+        sieve[0:2] = [False, False]
 
-def fake_reshape(
-    true_tensor: "torch.Tensor",  # noqa: F821
-    sh: Dict[int, Any],  # noqa: F821
-    fake_tensor: Optional["FakeTensor"] = None,  # noqa: F821
-    fake_mode: Optional["FakeTensorMode"] = None,  # noqa: F821
-) -> "FakeTensor":  # noqa: F821
-    """
-    Changes the shape of a true tensor to make it dynamic.
+        for i in range(2, int(n**0.5) + 1):
+            if sieve[i]:
+                # Élimine les multiples de i
+                sieve[i * i : n + 1 : i] = [False] * len(range(i * i, n + 1, i))
 
-    :param true_tensor: true tensor
-    :param sh: dynamic shape
-    :param fake_tensor: fake tensor, if None, make a fake one
-    :param fake_mode: fake tensor mode
-    :return: fake tensor
-    """
-    import torch
-
-    # deal with 0/1
-    for i in sh:
-        if true_tensor.shape[i] <= 1:
-            expanded_shape = list(true_tensor.shape)
-            expanded_shape[i] = _unique()
-            true_tensor = torch.empty(
-                tuple(expanded_shape), dtype=true_tensor.dtype, device=true_tensor.device
+        return [i for i, prime in enumerate(sieve) if prime and i >= 13]
+
+    def _unique(self) -> int:
+        i = 0
+        c = self._candidates[i]
+        while c in self._unique_ or c in self._mapping_int:
+            i += 1
+            assert i < len(
+                self._candidates
+            ), f"Two many unique dimensions to generate, requested: {len(self._unique_)}"
+            c = self._candidates[i]
+        self._unique_.add(c)
+        return c
+
+    def from_tensor(self, x, static_shapes=False) -> "FakeTensor":  # noqa: F821
+        """
+        Returns a fake tensor.
+        ``pytorch`` returns the same name for the same dimension.
+        """
+        fake = self.fake_mode.from_tensor(x, static_shapes=static_shapes)
+        for i, s in zip(x.shape, fake.shape):
+            assert i not in self._mapping_int or self._mapping_int[i] == s, (
+                f"Inconsistency between {x.shape} and {fake.shape}, "
+                f"mapping has {self._mapping_int[i]} and s={s}"
+            )
+            self._mapping_int[i] = s
+        return fake
+
+    def fake_reshape(
+        self,
+        true_tensor: "torch.Tensor",  # noqa: F821
+        sh: Dict[int, Any],  # noqa: F821
+        fake_tensor: Optional["FakeTensor"] = None,  # noqa: F821
+    ) -> "FakeTensor":  # noqa: F821
+        """
+        Changes the shape of a true tensor to make it dynamic.
+
+        :param true_tensor: true tensor
+        :param sh: dynamic shape
+        :param fake_tensor: fake tensor, if None, make a fake one
+        :return: fake tensor
+        """
+        import torch
+
+        # deal with 0/1
+        for i in sh:
+            if true_tensor.shape[i] <= 1:
+                expanded_shape = list(true_tensor.shape)
+                expanded_shape[i] = self._unique()
+                true_tensor = torch.empty(
+                    tuple(expanded_shape), dtype=true_tensor.dtype, device=true_tensor.device
+                )
+
+        # deal with equivalent dimension
+        new_shape = list(true_tensor.shape)
+        mapping = {}
+        for i, s in sh.items():
+            d = true_tensor.shape[i]
+            if d not in mapping:
+                mapping[d] = s
+            elif mapping[d] != s:
+                d = self._unique()
+                mapping[d] = s
+                new_shape[i] = d
+        true_tensor = torch.empty(
+            tuple(new_shape), dtype=true_tensor.dtype, device=true_tensor.device
+        )
+
+        # now switch to FakeTensor
+        fake_tensor = self.from_tensor(true_tensor, static_shapes=False)
+        new_shape = list(true_tensor.shape)
+        for i in sh:
+            new_shape[i] = fake_tensor.shape[i]
+
+        reduced_tensor = self.from_tensor(true_tensor, static_shapes=True).sum(
+            axis=tuple(sorted(sh)), keepdim=True
+        )
+        return reduced_tensor.expand(*new_shape)
+
+    def make_fake(self, x: Any) -> Optional["FakeTensor"]:  # noqa: F821
+        """See :func:`onnx_diagnostic.helpers.fake_tensor_helper.make_fake`."""
+        if x is None:
+            return None
+        if isinstance(x, (list, tuple)):
+            return x.__class__([self.make_fake(i) for i in x])
+        if isinstance(x, dict):
+            return {k: self.make_fake(v) for k, v in x.items()}
+        if x.__class__.__name__ in {"DynamicCache", "StaticCache", "HybridCache"}:
+            assert hasattr(x, "layers"), (
+                f"Une more recent version of transformers (>=4.55), "
+                f"'layers' not found in class {type(x)}"
+            )
+            for layer in x.layers:
+                assert hasattr(layer, "keys") and hasattr(layer, "values"), (
+                    f"Une more recent version of transformers (>=4.55), 'layers' "
+                    f"not found in class {type(layer)} ({dir(layer)})"
+                )
+                layer.keys = self.make_fake(layer.keys)
+                layer.values = self.make_fake(layer.values)
+            return x
+        if x.__class__.__name__ == "EncoderDecoderCache":
+            self.make_fake(x.self_attention_cache)
+            self.make_fake(x.cross_attention_cache)
+            return x
+        if hasattr(x, "shape"):
+            return self.from_tensor(x, static_shapes=False)
+        from . import string_type
+
+        raise TypeError(
+            f"Unexpected type {type(x)} for x, content is {string_type(x, with_shape=True)}"
+        )
+
+    def make_fake_with_dynamic_dimensions(self, x: Any, dynamic_shapes: Any) -> Any:
+        """
+        See
+        :func:`onnx_diagnostic.export.shape_helper.make_fake_with_dynamic_dimensions`.
+        """
+        if x is None:
+            return None, None
+        if isinstance(x, (list, tuple)):
+            return x.__class__(
+                [
+                    self.make_fake_with_dynamic_dimensions(i, dynamic_shapes=ds)
+                    for i, ds in zip(x, dynamic_shapes)
+                ]
+            )
+        if isinstance(x, dict):
+            return {
+                k: self.make_fake_with_dynamic_dimensions(v, dynamic_shapes=dynamic_shapes[k])
+                for k, v in x.items()
+            }
+        if x.__class__.__name__ in {"DynamicCache", "StaticCache", "HybridCache"}:
+            assert hasattr(x, "layers"), (
+                f"Une more recent version of transformers (>=4.55), "
+                f"'layers' not found in class {type(x)}"
+            )
+            assert isinstance(dynamic_shapes, list) and (
+                not dynamic_shapes or not isinstance(dynamic_shapes[0], list)
+            ), f"Unexpected dynamic_shapes={dynamic_shapes} for a DynamicCache"
+            for il, layer in enumerate(x.layers):
+                assert hasattr(layer, "keys") and hasattr(layer, "values"), (
+                    f"Une more recent version of transformers (>=4.55), 'layers' "
+                    f"not found in class {type(layer)} ({dir(layer)})"
+                )
+                layer.keys = self.make_fake_with_dynamic_dimensions(
+                    layer.keys, dynamic_shapes=dynamic_shapes[il * 2]
+                )
+                layer.values = self.make_fake_with_dynamic_dimensions(
+                    layer.values, dynamic_shapes=dynamic_shapes[il * 2 + 1]
+                )
+            return x
+        if x.__class__.__name__ == "EncoderDecoderCache":
+            self.make_fake_with_dynamic_dimensions(
+                x.self_attention_cache, dynamic_shapes=dynamic_shapes[0]
             )
+            self.make_fake_with_dynamic_dimensions(
+                x.cross_attention_cache, dynamic_shapes=dynamic_shapes[1]
+            )
+            return x
+        if hasattr(x, "shape"):
+            assert dynamic_shapes is None or isinstance(dynamic_shapes, dict), (
+                f"dynamic_shapes must be a dictionary at this stage but "
+                f"dynamic_shapes={dynamic_shapes}"
+            )
+            # We need to overwrite the values.
+            new_shape = []
+            for idim, dim in enumerate(x.shape):
+                if dynamic_shapes is not None and idim in dynamic_shapes:
+                    s = dynamic_shapes[idim]
+                    assert isinstance(s, str), (
+                        f"Unexpected type {type(s)} in dynamic_shapes={dynamic_shapes} "
+                        f"at index {idim}"
+                    )
+                    if s in self._mapping_str:
+                        dim = self._mapping_str[s]
+                    else:
+                        i = self._unique()
+                        self._mapping_str[s] = i
+                        dim = i
+                assert isinstance(dim, int), (
+                    f"Unexpected type {type(dim)}, dynamic_shapes={dynamic_shapes} "
+                    f"at index {idim}, dim={dim}"
+                )
+                new_shape.append(dim)
+            if tuple(new_shape) != x.shape:
+                import torch
 
-    # deal with equivalent dimension
-    new_shape = list(true_tensor.shape)
-    mapping = {}
-    for i, s in sh.items():
-        d = true_tensor.shape[i]
-        if d not in mapping:
-            mapping[d] = s
-        elif mapping[d] != s:
-            d = _unique()
-            mapping[d] = s
-            new_shape[i] = d
-    true_tensor = torch.empty(
-        tuple(new_shape), dtype=true_tensor.dtype, device=true_tensor.device
-    )
-
-    # now switch to FakeTensor
-    if fake_mode is None:
-        from torch.fx.experimental.symbolic_shapes import ShapeEnv
-        from torch._subclasses.fake_tensor import FakeTensorMode
-
-        shape_env = ShapeEnv()
-        fake_mode = FakeTensorMode(shape_env=shape_env)
-    if fake_tensor is None:
-        fake_tensor = fake_mode.from_tensor(true_tensor, static_shapes=False)
-    assert fake_mode is not None, "fake_mode must be provided"
-
-    new_shape = list(true_tensor.shape)
-    for i in sh:
-        new_shape[i] = fake_tensor.shape[i]
-
-    reduced_tensor = fake_mode.from_tensor(true_tensor, static_shapes=True).sum(
-        axis=tuple(sorted(sh)), keepdim=True
-    )
-    return reduced_tensor.expand(*new_shape)
+                x = torch.empty(tuple(new_shape), dtype=x.dtype, device=x.device)
+
+            t = self.fake_reshape(x, dynamic_shapes)  # type: ignore[arg-type]
+            assert t.device == x.device, f"device mismatch {x.device} -> {t.device}"
+            assert t.dtype == x.dtype, f"dtype mismatch {x.dtype} -> {t.dtype}"
+            return t
+        from ..helpers import string_type
+
+        raise TypeError(
+            f"Unexpected type {type(x)} for x, content is {string_type(x, with_shape=True)}"
+        )
 
 
 def make_fake(
-    x: Any, fake_mode: Optional["FakeTensorMode"] = None  # noqa: F821
-) -> Tuple[Optional["FakeTensor"], Optional["FakeTensorMode"]]:  # noqa: F821
+    x: Any, context: Optional[FakeTensorContext] = None
+) -> Tuple[Optional["FakeTensor"], Optional[FakeTensorContext]]:  # noqa: F821
     """
     Replaces all tensors by fake tensors.
     This modification happens inplace for caches.
@@ -114,40 +268,6 @@ def make_fake(
     """
     if x is None:
         return None, None
-    if fake_mode is None:
-        from torch.fx.experimental.symbolic_shapes import ShapeEnv
-        from torch._subclasses.fake_tensor import FakeTensorMode
-
-        shape_env = ShapeEnv()
-        fake_mode = FakeTensorMode(shape_env=shape_env)
-
-    if isinstance(x, (list, tuple)):
-        return x.__class__([make_fake(i, fake_mode=fake_mode)[0] for i in x]), fake_mode
-    if isinstance(x, dict):
-        return {k: make_fake(v, fake_mode=fake_mode)[0] for k, v in x.items()}, fake_mode
-
-    if x.__class__.__name__ in {"DynamicCache", "StaticCache", "HybridCache"}:
-        assert hasattr(x, "layers"), (
-            f"Une more recent version of transformers (>=4.55), "
-            f"'layers' not found in class {type(x)}"
-        )
-        for layer in x.layers:
-            assert hasattr(layer, "keys") and hasattr(layer, "values"), (
-                f"Une more recent version of transformers (>=4.55), 'layers' "
-                f"not found in class {type(layer)} ({dir(layer)})"
-            )
-            layer.keys = make_fake(layer.keys, fake_mode=fake_mode)[0]
-            layer.values = make_fake(layer.values, fake_mode=fake_mode)[0]
-        return x, fake_mode
-    if x.__class__.__name__ == "EncoderDecoderCache":
-        make_fake(x.self_attention_cache, fake_mode=fake_mode)
-        make_fake(x.cross_attention_cache, fake_mode=fake_mode)
-        return x, fake_mode
-    if hasattr(x, "shape"):
-        t = fake_mode.from_tensor(x, static_shapes=False)
-        return t, fake_mode
-    from . import string_type
-
-    raise TypeError(
-        f"Unexpected type {type(x)} for x, content is {string_type(x, with_shape=True)}"
-    )
+    if context is None:
+        context = FakeTensorContext()
+    return context.make_fake(x), context

onnx_diagnostic/helpers/helper.py

@@ -1,6 +1,7 @@
 import ast
 import enum
 import inspect
+import itertools
 from dataclasses import is_dataclass, fields
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 import numpy as np
@@ -948,8 +949,8 @@ def flatten_object(x: Any, drop_keys: bool = False) -> Any:
         from .cache_helper import CacheKeyValue
 
         kc = CacheKeyValue(x)
-        res = flatten_object(kc.key_cache) + flatten_object(kc.value_cache)
-        return tuple(res)
+        return list(itertools.chain.from_iterable(zip(kc.key_cache, kc.value_cache)))
+
     if x.__class__.__name__ == "EncoderDecoderCache":
         res = flatten_object(x.self_attention_cache) + flatten_object(x.cross_attention_cache)
         return tuple(res)
@@ -1015,6 +1016,8 @@ def max_diff(
 
     You may use :func:`string_diff` to display the discrepancies in one string.
     """
+    if verbose >= 10:
+        print(f"[max_diff] {type(expected)} ? {type(got)}")
     if expected is None and got is None:
         return dict(abs=0, rel=0, sum=0, n=0, dnan=0)
 
@@ -1056,6 +1059,27 @@ def max_diff(
             allow_unique_tensor_with_list_of_one_element=False,
             hist=hist,
         )
+
+    if expected.__class__.__name__ == "CausalLMOutputWithPast":
+        if verbose >= 6:
+            print(
+                f"[max_diff] CausalLMOutputWithPast: {string_type(expected, with_shape=True)} "
+                f"? {string_type(got, with_shape=True)}"
+            )
+        if got.__class__.__name__ == "CausalLMOutputWithPast":
+            return max_diff(
+                [expected.logits, *flatten_object(expected.past_key_values)],
+                [got.logits, *flatten_object(got.past_key_values)],
+                debug_info=_debug(expected.__class__.__name__),
+                **_dkws,
+            )
+        return max_diff(
+            [expected.logits, *flatten_object(expected.past_key_values)],
+            got,
+            debug_info=_debug(expected.__class__.__name__),
+            **_dkws,
+        )
+
     if hasattr(expected, "to_tuple"):
         if verbose >= 6:
             print(f"[max_diff] to_tuple1: {string_type(expected)} ? {string_type(got)}")
@@ -1066,36 +1090,6 @@ def max_diff(
             print(f"[max_diff] to_tuple2: {string_type(expected)} ? {string_type(got)}")
         return max_diff(expected, got.to_tuple(), debug_info=_debug("to_tuple2"), **_dkws)
 
-        if isinstance(got, (list, tuple)):
-            if len(got) != 1:
-                if verbose >= 6:
-                    print(
-                        f"[max_diff] list,tuple,2: {string_type(expected)} "
-                        f"? {string_type(got)}"
-                    )
-                if verbose > 2:
-                    import torch
-
-                    print(
-                        f"[max_diff] (a) inf because len(expected)={len(expected)}!=1, "
-                        f"len(got)={len(got)}, level={level}, _index={_index}"
-                    )
-                    for i, (a, b) in enumerate(zip(expected, got)):
-                        if isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor):
-                            print(
-                                f"    i={i} expected {a.dtype}:{a.shape}, "
-                                f"has {b.dtype}:{b.shape}, _index={_index}"
-                            )
-                        else:
-                            print(
-                                f"    i={i} a is {type(a)}, "
-                                f"b is {type(b)}, _index={_index}"
-                            )
-                return dict(abs=np.inf, rel=np.inf, sum=np.inf, n=np.inf, dnan=np.inf)
-            if verbose >= 6:
-                print(f"[max_diff] list,tuple,1: {string_type(expected)} ? {string_type(got)}")
-            return max_diff(expected, got[0], debug_info=_debug("lt1"), **_dkws)
-
     if isinstance(expected, (tuple, list)):
         if verbose >= 6:
             print(f"[max_diff] list,tuple,0: {string_type(expected)} ? {string_type(got)}")
@@ -1484,7 +1478,7 @@ def max_diff(
             return dict(abs=np.inf, rel=np.inf, sum=np.inf, n=np.inf, dnan=np.inf)
         if verbose >= 6:
             print(
-                f"[max_diff] {expected.__class__.__name__}: "
+                f"[max_diff*] {expected.__class__.__name__}: "
                 f"{string_type(expected)} ? {string_type(got)}"
             )
         expected_args, _spec = torch.utils._pytree.tree_flatten(expected)