fix make_fake

Author: xadupre

_unittests/ut_export/test_shape_helper.py

@@ -183,6 +183,40 @@ def test_make_fake_with_dynamic_dimensions_tensor(self):
         self.assertEqual(reshaped.shape[3], 96)
         self.assertNotEqual(reshaped.shape[0], reshaped.shape[2])
 
+    @requires_transformers("4.55")
+    @requires_torch("2.9")
+    def test_make_fake_with_dynamic_dimensions_two_tensors(self):
+        res = make_fake_with_dynamic_dimensions(
+            (
+                torch.rand((2, 32, 30, 96), dtype=torch.float16),
+                torch.rand((2, 32, 30, 96), dtype=torch.float16),
+            ),
+            ({0: "batch", 2: "cache_length"}, {0: "batch", 2: "cache_length"}),
+        )
+        reshaped = res[0][0]
+        self.assertIsInstance(reshaped.shape[0], torch.SymInt)
+        self.assertIsInstance(reshaped.shape[2], torch.SymInt)
+        self.assertEqual(reshaped.shape[1], 32)
+        self.assertEqual(reshaped.shape[3], 96)
+        self.assertNotEqual(reshaped.shape[0], reshaped.shape[2])
+        self.assertEqual(str(res[0][0].shape), str(res[0][1].shape))
+        sh1 = res[0][0].shape
+        sh2 = res[0][1].shape
+        self.assertEqual(sh1[0], sh2[0])
+        self.assertEqual(sh1[1], sh2[1])
+        self.assertEqual(sh1[2], sh2[2])
+        self.assertEqual(sh1[3], sh2[3])
+
+    def test_make_fake_with_dynamic_dimensions_attention(self):
+        query = torch.rand((1, 2, 1, 96), dtype=torch.float32)
+        key = torch.rand((1, 2, 4, 96), dtype=torch.float32)
+        value = torch.rand((1, 2, 4, 96), dtype=torch.float32)
+        ds = ({0: "batch", 2: "seq1"}, {0: "batch", 2: "seq2"}, {0: "batch", 2: "seq2"})
+        fake_inputs, _ = make_fake_with_dynamic_dimensions((query, key, value), ds)
+        self.assertEqual(fake_inputs[1].shape, fake_inputs[2].shape)
+        self.assertEqual(fake_inputs[0].shape[0], fake_inputs[1].shape[0])
+        self.assertEqual(fake_inputs[0].shape[0], fake_inputs[2].shape[0])
+
     @requires_transformers("4.55")
     @requires_torch("2.9")
     def test_make_fake_with_dynamic_dimensions_whole(self):

_unittests/ut_helpers/test_fake_tensor_helper.py

@@ -3,72 +3,72 @@
 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, fake_reshape
+from onnx_diagnostic.helpers.fake_tensor_helper import make_fake, FakeTensorContext
 
 
 class TestMakeTensorHelper(ExtTestCase):
 
+    @requires_transformers("4.55")
+    def test_fake_inputs(self):
+        inputs, _ = make_fake(
+            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),
+                        ),
+                    ]
+                ),
+            )
+        )
+        flat = flatten_object(inputs, drop_keys=True)
+        for t in flat:
+            self.assertIsInstance(t, torch.Tensor)
+            assert all(
+                isinstance(s, torch.SymInt) for s in t.shape
+            ), f"Wrong type {[type(s) for s in t.shape]} in {t.shape}"
+
     def test_fake_reshape_generic(self):
         t = torch.zeros((2, 3, 4, 5), dtype=torch.float32)
-        reshaped = fake_reshape(t, {0: "batch", 2: "seq_length"})
+        reshaped = FakeTensorContext().fake_reshape(t, {0: "batch", 2: "seq_length"})
         self.assertIsInstance(reshaped.shape[0], torch.SymInt)
         self.assertIsInstance(reshaped.shape[2], torch.SymInt)
         self.assertEqual(reshaped.shape[1], 3)
         self.assertEqual(reshaped.shape[3], 5)
 
     def test_fake_reshape_dim_1(self):
         t = torch.zeros((1, 3, 4, 5), dtype=torch.float32)
-        reshaped = fake_reshape(t, {0: "batch", 2: "seq_length"})
+        reshaped = FakeTensorContext().fake_reshape(t, {0: "batch", 2: "seq_length"})
         self.assertIsInstance(reshaped.shape[0], torch.SymInt)
         self.assertIsInstance(reshaped.shape[2], torch.SymInt)
         self.assertEqual(reshaped.shape[1], 3)
         self.assertEqual(reshaped.shape[3], 5)
 
     def test_fake_reshape_dim_0(self):
         t = torch.zeros((0, 3, 4, 5), dtype=torch.float32)
-        reshaped = fake_reshape(t, {0: "batch", 2: "seq_length"})
+        reshaped = FakeTensorContext().fake_reshape(t, {0: "batch", 2: "seq_length"})
         self.assertIsInstance(reshaped.shape[0], torch.SymInt)
         self.assertIsInstance(reshaped.shape[2], torch.SymInt)
         self.assertEqual(reshaped.shape[1], 3)
         self.assertEqual(reshaped.shape[3], 5)
 
     def test_fake_reshape_different(self):
         t = torch.zeros((2, 3, 2, 5), dtype=torch.float32)
-        reshaped = fake_reshape(t, {0: "batch", 2: "seq_length"})
+        reshaped = FakeTensorContext().fake_reshape(t, {0: "batch", 2: "seq_length"})
         self.assertIsInstance(reshaped.shape[0], torch.SymInt)
         self.assertIsInstance(reshaped.shape[2], torch.SymInt)
         self.assertEqual(reshaped.shape[1], 3)
         self.assertEqual(reshaped.shape[3], 5)
         self.assertNotEqual(reshaped.shape[0], reshaped.shape[2])
 
-    @requires_transformers("4.55")
-    def test_fake_inputs(self):
-        inputs, _ = make_fake(
-            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),
-                        ),
-                    ]
-                ),
-            )
-        )
-        flat = flatten_object(inputs, drop_keys=True)
-        for t in flat:
-            self.assertIsInstance(t, torch.Tensor)
-            assert all(
-                isinstance(s, torch.SymInt) for s in t.shape
-            ), f"Wrong type {[type(s) for s in t.shape]} in {t.shape}"
-
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

_unittests/ut_torch_export_patches/test_patch_transformers.py

@@ -168,9 +168,7 @@ def forward(self, query, key, value):
         # dynamic
         ds = ({0: "batch", 2: "seq1"}, {0: "batch", 2: "seq2"}, {0: "batch", 2: "seq2"})
         fake_inputs, _ = make_fake_with_dynamic_dimensions((query, key, value), ds)
-        print("****", fake_inputs)
-        epd = torch.export.export(model, fake_inputs)  # , dynamic_shapes=use_dyn_not_str(ds))
-        print(epq)
+        epd = torch.export.export(model, fake_inputs, dynamic_shapes=use_dyn_not_str(ds))
         got = epd.module()(query, key, value)
         self.assertEqualArray(expected, got)
 

onnx_diagnostic/export/shape_helper.py

@@ -1,6 +1,5 @@
 from typing import Any, Dict, List, Set, Optional, Tuple, Union
 from ..helpers.cache_helper import flatten_unflatten_for_dynamic_shapes
-from ..helpers.fake_tensor_helper import fake_reshape
 from .dynamic_shapes import ModelInputs
 
 
@@ -203,14 +202,14 @@ def guess_dynamic_shapes_from_inputs(
 
 
 def make_fake_with_dynamic_dimensions(
-    x: Any,
-    dynamic_shapes: Any,
-    fake_mode: Optional["FakeTensorMode"] = None,  # noqa: F821
-) -> Tuple[Any, "FakeTensorMode"]:  # noqa: F821
+    x: Any, dynamic_shapes: Any, context: Optional["FakeTensorContext"] = None  # noqa: F821
+) -> Tuple[Any, "FakeTensorContext"]:  # 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`.
+    Parameter ``existing`` is used to reused the same object when the dynamic
+    dimension is given the same name as another one.
 
     A simple tensor:
 
@@ -227,6 +226,24 @@ def make_fake_with_dynamic_dimensions(
         )
         print(inputs)
 
+    Two tensors:
+
+    .. runpython::
+        :showcode:
+
+        import torch
+        from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
+        from onnx_diagnostic.export.shape_helper import make_fake_with_dynamic_dimensions
+
+        inputs, _ = make_fake_with_dynamic_dimensions(
+            (
+                torch.rand((2, 3, 4, 5), dtype=torch.float32),
+                torch.rand((2, 3, 4, 5), dtype=torch.float32),
+            ),
+            ({0: "batch", 2: "cache_length"}, {0: "batch", 2: "cache_length"}),
+        )
+        print(inputs)
+
     With a cache:
 
     .. runpython::
@@ -271,68 +288,9 @@ def make_fake_with_dynamic_dimensions(
     """
     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
+    if context is None:
+        from ..helpers.fake_tensor_helper import FakeTensorContext
 
-        shape_env = ShapeEnv()
-        fake_mode = FakeTensorMode(shape_env=shape_env)
+        context = FakeTensorContext()
 
-    if isinstance(x, (list, tuple)):
-        return (
-            x.__class__(
-                [
-                    make_fake_with_dynamic_dimensions(
-                        i, fake_mode=fake_mode, dynamic_shapes=ds
-                    )[0]
-                    for i, ds in zip(x, dynamic_shapes)
-                ]
-            ),
-            fake_mode,
-        )
-    if isinstance(x, dict):
-        return {
-            k: make_fake_with_dynamic_dimensions(
-                v, fake_mode=fake_mode, dynamic_shapes=dynamic_shapes[k]
-            )[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)}"
-        )
-        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 = make_fake_with_dynamic_dimensions(
-                layer.keys, fake_mode=fake_mode, dynamic_shapes=dynamic_shapes[il * 2]
-            )[0]
-            layer.values = make_fake_with_dynamic_dimensions(
-                layer.values, fake_mode=fake_mode, dynamic_shapes=dynamic_shapes[il * 2 + 1]
-            )[0]
-        return x, fake_mode
-    if x.__class__.__name__ == "EncoderDecoderCache":
-        make_fake_with_dynamic_dimensions(
-            x.self_attention_cache, fake_mode=fake_mode, dynamic_shapes=dynamic_shapes[0]
-        )
-        make_fake_with_dynamic_dimensions(
-            x.cross_attention_cache, fake_mode=fake_mode, dynamic_shapes=dynamic_shapes[1]
-        )
-        return x, fake_mode
-    if hasattr(x, "shape"):
-        t = fake_reshape(x, dynamic_shapes, fake_mode=fake_mode)
-        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, fake_mode
-    from ..helpers import string_type
-
-    raise TypeError(
-        f"Unexpected type {type(x)} for x, content is {string_type(x, with_shape=True)}"
-    )
+    return context.make_fake_with_dynamic_dimensions(x, dynamic_shapes), context