lint

Author: xadupre

CHANGELOGS.rst

@@ -4,6 +4,7 @@ Change Logs
 0.7.16
 ++++++
 
+* :pr:`272`: makes patches woth with FakeTensor
 * :pr:`270`: add export sample code to export a specific model id with the appropriate inputs
 * :pr:`269`: adds one unit test to track a patch fixing broadcast output shape
 * :pr:`267`: patches ``sdpa_attention_forward`` because of a control flow (``transformers>=5.0``)

_unittests/ut_helpers/test_fake_tensor_helper.py

@@ -1,12 +1,13 @@
 import unittest
 import torch
-from onnx_diagnostic.ext_test_case import ExtTestCase
+from onnx_diagnostic.ext_test_case import ExtTestCase, requires_transformers
 from onnx_diagnostic.helpers import flatten_object
 from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
 from onnx_diagnostic.helpers.fake_tensor_helper import make_fake
 
 
 class TestMakeTensorHelper(ExtTestCase):
+    @requires_transformers("4.55")
     def test_fake_inputs(self):
         inputs, _ = make_fake(
             dict(

onnx_diagnostic/_command_lines_parser.py

@@ -1128,7 +1128,7 @@ def main(argv: Optional[List[Any]] = None):
                 raise ValueError(
                     f"Unknown command {cmd!r}, it should be in {list(sorted(parsers))}."
                 )
-            parser = parsers[cmd]()
+            parser = parsers[cmd]()  # type: ignore[operator]
             parser.parse_args(argv[1:])
         raise RuntimeError("The programme should have exited before.")
 

onnx_diagnostic/export/dynamic_shapes.py

@@ -2,8 +2,9 @@
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 import numpy as np
 import torch
-from ..helpers import string_type
+from ..helpers import string_type, flatten_object
 from ..helpers.cache_helper import flatten_unflatten_for_dynamic_shapes
+from ..helpers.fake_tensor_helper import make_fake
 
 DYNAMIC_SHAPES = Tuple[Tuple[Any, ...], Dict[str, Any]]
 
@@ -32,6 +33,51 @@ def _flat_list(li: List[Any]) -> List[Dict[int, str]]:
     return res
 
 
+def make_fake_with_dynamic_dimensions(
+    x: Optional[Any],
+    dynamic_shapes: Any,
+    fake_mode: Optional["FakeTensorMode"] = None,  # noqa: F821
+) -> Optional[Tuple["FakeTensor", "FaleTensorMode"]]:  # noqa: F821
+    """
+    Replaces all tensors by fake tensor respecting the same
+    constraints as the following dynamic shapes.
+    This uses function :func:`onnx_diagnostic.helpers.fake_tensor_helper.make_fake`.
+
+    .. runpython::
+        :showcode:
+
+        from onnx_diagnostic.export.dynamic_shapes import make_fake_with_dynamic_dimensions
+
+        inputs, _ = make_fake_with_dynamic_dimensions(
+            dict(
+                input_ids=torch.randint(30360, size=(2, 3), dtype=torch.int64),
+                attention_mask=torch.randint(1, size=(2, 33), dtype=torch.int64),
+                position_ids=torch.randint(32, size=(2, 3), dtype=torch.int64),
+                past_key_values=make_dynamic_cache(
+                    [
+                        (
+                            torch.rand((2, 32, 30, 96), dtype=torch.float16),
+                            torch.rand((2, 32, 30, 96), dtype=torch.float16),
+                        ),
+                        (
+                            torch.rand((2, 32, 30, 96), dtype=torch.float16),
+                            torch.rand((2, 32, 30, 96), dtype=torch.float16),
+                        ),
+                    ]
+                ),
+            )
+        )
+        print(inputs)
+    """
+    fake_inputs = make_fake(x, fake_mode=fake_mode)
+    flat_inputs = flatten_object(fake_inputs, drop_keys=True)
+    flat_ds = flatten_dynamic_shapes(dynamic_shapes)
+    assert len(flat_inputs) == len(flat_ds), (
+        f"Mismatch between the number of input tensor {len(flat_inputs)} "
+        f"and the number of dynamic_shapes {len(flat_ds)}"
+    )
+
+
 class CoupleInputsDynamicShapes:
     """
     Pair inputs / dynamic shapes.

onnx_diagnostic/helpers/fake_tensor_helper.py

@@ -2,11 +2,13 @@
 
 
 def make_fake(
-    x: Any, fake_mode: Optional["FakeTensorMode"] = None  # noqa: F821
-) -> Tuple["FakeTensor", "FaleTensorMode"]:  # noqa: F821
+    x: Optional[Any], fake_mode: Optional["FakeTensorMode"] = None  # noqa: F821
+) -> Optional[Tuple["FakeTensor", "FaleTensorMode"]]:  # noqa: F821
     """
     Replaces all tensors by fake tensors.
     This modification happens inplace for caches.
+    This function is only implemented for cache with
+    ``transformers>=4.55``.
 
     .. runpython::
         :showcode:
@@ -49,12 +51,13 @@ def make_fake(
         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, 'layers' not found in class {type(x)}"
+        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, 'layers' "
+                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]

pyproject.toml

@@ -58,6 +58,10 @@ disable_error_code = ["name-defined"]
 module = ["onnx_diagnostic.helpers.helper"]
 disable_error_code = ["arg-type", "assignment", "attr-defined", "call-overload", "misc", "name-defined", "union-attr"]
 
+[[tool.mypy.overrides]]
+module = ["onnx_diagnostic.helpers.fake_tensor_helper"]
+disable_error_code = ["name-defined"]
+
 [[tool.mypy.overrides]]
 module = ["onnx_diagnostic.helpers.model_builder_helper"]
 disable_error_code = ["attr-defined", "import-untyped", "name-defined", "union-attr", "var-annotated"]