Makes patches work with FakeTensor (#272)

* fix patches

* fix patch

* fixes shape information

* fix issues

* add fake

* lint

* more about fake helper

* g

* mypy

* pypi

* fix shape

* mypy

Author: xadupre

CHANGELOGS.rst

@@ -4,6 +4,7 @@ Change Logs
 0.7.16
 ++++++
 
+* :pr:`272`: makes patches work 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``)
@@ -100,7 +101,7 @@ Change Logs
 +++++
 
 * :pr:`178`: add a patch for eager_mask to handle ``assert len(flat_dynamic_shapes) == num_placeholders - num_lifted_inputs``
-* :pr:`177`: changes for the next version of onnx, fixes all_dynamic_shape_from_inputs
+* :pr:`177`: changes for the next version of onnx, fixes all_dynamic_shapes_from_inputs
 
 0.7.3
 +++++

README.rst

@@ -43,9 +43,9 @@ You need then to remove those which are not dynamic in your model.
 
 .. code-block:: python
 
-  from onnx_diagnostic.export.shape_helper import all_dynamic_shape_from_inputs
+  from onnx_diagnostic.export.shape_helper import all_dynamic_shapes_from_inputs
 
-  dynamic_shapes = all_dynamic_shape_from_inputs(cache)
+  dynamic_shapes = all_dynamic_shapes_from_inputs(cache)
 
 It also implements tools to investigate, validate exported models (ExportedProgramm, ONNXProgram, ...).
 See `documentation of onnx-diagnostic <https://sdpython.github.io/doc/onnx-diagnostic/dev/>`_ and
@@ -109,13 +109,13 @@ Snapshot of usefuls tools
       ep = torch.export.export(model, args, kwargs=kwargs, dynamic_shapes=dynamic_shapes)
       # ...
 
-**all_dynamic_shape_from_inputs**
+**all_dynamic_shapes_from_inputs**
 
 .. code-block:: python
 
-  from onnx_diagnostic.export.shape_helper import all_dynamic_shape_from_inputs
+  from onnx_diagnostic.export.shape_helper import all_dynamic_shapes_from_inputs
 
-  dynamic_shapes = all_dynamic_shape_from_inputs(cache)
+  dynamic_shapes = all_dynamic_shapes_from_inputs(cache)
 
 **torch_export_rewrite**
 

_doc/api/helpers/fake_tensor_helper.rst

@@ -0,0 +1,7 @@
+
+onnx_diagnostic.helpers.fake_tensor_helper
+==========================================
+
+.. automodule:: onnx_diagnostic.helpers.fake_tensor_helper
+    :members:
+    :no-undoc-members:

_doc/api/helpers/index.rst

@@ -11,6 +11,7 @@ onnx_diagnostic.helpers
     cache_helper
     config_helper
     doc_helper
+    fake_tensor_helper
     graph_helper
     helper
     _log_helper

_doc/index.rst

@@ -33,15 +33,15 @@ Patches can be enabled as follows with function
       # ...
 
 Dynamic shapes are difficult to guess for caches, function
-:func:`onnx_diagnostic.export.shape_helper.all_dynamic_shape_from_inputs`
+:func:`onnx_diagnostic.export.shape_helper.all_dynamic_shapes_from_inputs`
 returns a structure defining all dimensions as dynamic.
 You need then to remove those which are not dynamic in your model.
 
 .. code-block:: python
 
-  from onnx_diagnostic.export.shape_helper import all_dynamic_shape_from_inputs
+  from onnx_diagnostic.export.shape_helper import all_dynamic_shapes_from_inputs
 
-  dynamic_shapes = all_dynamic_shape_from_inputs(cache)
+  dynamic_shapes = all_dynamic_shapes_from_inputs(cache)
 
 It also implements tools to investigate, validate exported models (ExportedProgramm, ONNXProgram, ...).
 :func:`onnx_diagnostic.torch_export_patches.torch_export_patches`.
@@ -134,16 +134,16 @@ See :func:`onnx_diagnostic.torch_export_patches.torch_export_rewrite`.
       ep = torch.export.export(model, args, kwargs=kwargs, dynamic_shapes=dynamic_shapes)
       # ...
 
-all_dynamic_shape_from_inputs
+all_dynamic_shapes_from_inputs
 +++++++++++++++++++++++++++++
 
-See :func:`onnx_diagnostic.export.shape_helper.all_dynamic_shape_from_inputs`.
+See :func:`onnx_diagnostic.export.shape_helper.all_dynamic_shapes_from_inputs`.
 
 .. code-block:: python
 
-  from onnx_diagnostic.export.shape_helper import all_dynamic_shape_from_inputs
+  from onnx_diagnostic.export.shape_helper import all_dynamic_shapes_from_inputs
 
-  dynamic_shapes = all_dynamic_shape_from_inputs(cache)
+  dynamic_shapes = all_dynamic_shapes_from_inputs(cache)
 
 string_type
 +++++++++++

_doc/recipes/plot_dynamic_shapes_json.py

@@ -21,7 +21,7 @@
 from onnx_diagnostic import doc
 from onnx_diagnostic.helpers import string_type
 from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
-from onnx_diagnostic.export.shape_helper import all_dynamic_shape_from_inputs
+from onnx_diagnostic.export.shape_helper import all_dynamic_shapes_from_inputs
 
 bsize, nheads, slen, dim = 2, 1, 30, 96
 
@@ -39,9 +39,9 @@
 print(string_type(inputs, with_shape=True))
 
 # %%
-# Function :func:`onnx_diagnostic.export.shape_helper.all_dynamic_shape_from_inputs`
+# Function :func:`onnx_diagnostic.export.shape_helper.all_dynamic_shapes_from_inputs`
 # produces the corresponding dynamic shapes assuming they are all dynamic.
-ds = all_dynamic_shape_from_inputs(inputs)
+ds = all_dynamic_shapes_from_inputs(inputs)
 pprint.pprint(ds)
 
 # %%

_doc/recipes/plot_dynamic_shapes_what.py

@@ -16,7 +16,7 @@
 from onnx_diagnostic import doc
 from onnx_diagnostic.helpers import string_type
 from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
-from onnx_diagnostic.export.shape_helper import all_dynamic_shape_from_inputs
+from onnx_diagnostic.export.shape_helper import all_dynamic_shapes_from_inputs
 from onnx_diagnostic.torch_models.hghub import get_untrained_model_with_inputs
 from onnx_diagnostic.torch_export_patches import torch_export_patches
 
@@ -34,9 +34,9 @@
 print(string_type(inputs, with_shape=True))
 
 # %%
-# Function :func:`onnx_diagnostic.export.shape_helper.all_dynamic_shape_from_inputs`
+# Function :func:`onnx_diagnostic.export.shape_helper.all_dynamic_shapes_from_inputs`
 # produces the corresponding dynamic shapes assuming they are all dynamic.
-ds = all_dynamic_shape_from_inputs(inputs)
+ds = all_dynamic_shapes_from_inputs(inputs)
 pprint.pprint(ds)
 
 # %%
@@ -56,13 +56,13 @@
 
 # %%
 # And the input shapes.
-ds = all_dynamic_shape_from_inputs(inputs)
+ds = all_dynamic_shapes_from_inputs(inputs)
 if ds["past_key_values"]:
     print("transformers implemented serialization function for StaticCache.")
 else:
     print("We need to use serialization function implemented in this package.")
     with torch_export_patches(patch_transformers=True):
-        ds = all_dynamic_shape_from_inputs(inputs)
+        ds = all_dynamic_shapes_from_inputs(inputs)
 
 # %%
 # That gives.

_unittests/ut_export/test_shape_helper.py

@@ -1,27 +1,28 @@
 import unittest
 import torch
 from onnx_diagnostic.ext_test_case import ExtTestCase, requires_transformers, requires_torch
-from onnx_diagnostic.export.shape_helper import (
-    all_dynamic_shape_from_inputs,
-    guess_dynamic_shapes_from_inputs,
-)
+from onnx_diagnostic.torch_models.hghub import get_untrained_model_with_inputs
+from onnx_diagnostic.torch_export_patches import torch_export_patches
+from onnx_diagnostic.helpers import flatten_object
 from onnx_diagnostic.helpers.cache_helper import (
     make_dynamic_cache,
     make_sliding_window_cache,
     make_encoder_decoder_cache,
     make_static_cache,
 )
-from onnx_diagnostic.torch_models.hghub import get_untrained_model_with_inputs
-from onnx_diagnostic.torch_export_patches import torch_export_patches
+from onnx_diagnostic.export.shape_helper import (
+    all_dynamic_shapes_from_inputs,
+    guess_dynamic_shapes_from_inputs,
+    make_fake_with_dynamic_dimensions,
+)
 
 
 class TestShapeHelper(ExtTestCase):
-
     @requires_transformers("4.52")
     @requires_torch("2.7.99")
     def test_all_dynamic_shape_from_cache(self):
         cache = make_dynamic_cache([(torch.ones((2, 2)), (torch.ones((2, 2)) * 2))])
-        ds = all_dynamic_shape_from_inputs(cache)
+        ds = all_dynamic_shapes_from_inputs(cache)
         self.assertEqual([[{0: "d_0_0", 1: "d_0_1"}], [{0: "d_1_0", 1: "d_1_1"}]], ds)
 
     @requires_torch("2.7.99")
@@ -122,17 +123,17 @@ def test_all_dynamic_shape_all_transformers_cache(self):
         with torch_export_patches(patch_transformers=True):
             for cache, exds in caches:
                 with self.subTest(cache_name=cache.__class__.__name__):
-                    ds = all_dynamic_shape_from_inputs(cache)
+                    ds = all_dynamic_shapes_from_inputs(cache)
                     self.assertEqual(exds, ds)
 
     @requires_transformers("4.52")
     @requires_torch("2.7.99")
-    def test_all_dynamic_shape_from_inputs(self):
-        ds = all_dynamic_shape_from_inputs((torch.randn((5, 6)), torch.randn((1, 6))))
+    def test_all_dynamic_shapes_from_inputs(self):
+        ds = all_dynamic_shapes_from_inputs((torch.randn((5, 6)), torch.randn((1, 6))))
         self.assertEqual(({0: "d_0_0", 1: "d_0_1"}, {0: "d_1_0", 1: "d_1_1"}), ds)
-        ds = all_dynamic_shape_from_inputs([torch.randn((5, 6)), torch.randn((1, 6))])
+        ds = all_dynamic_shapes_from_inputs([torch.randn((5, 6)), torch.randn((1, 6))])
         self.assertEqual([{0: "d_0_0", 1: "d_0_1"}, {0: "d_1_0", 1: "d_1_1"}], ds)
-        ds = all_dynamic_shape_from_inputs(
+        ds = all_dynamic_shapes_from_inputs(
             (torch.randn((5, 6)), torch.randn((1, 6))), dim_prefix=torch.export.Dim.AUTO
         )
         self.assertEqual(
@@ -145,9 +146,9 @@ def test_all_dynamic_shape_from_inputs(self):
 
     @requires_transformers("4.52")
     @requires_torch("2.7.99")
-    def test_all_dynamic_shape_from_inputs_dynamic_cache(self):
+    def test_all_dynamic_shapes_from_inputs_dynamic_cache(self):
         data = get_untrained_model_with_inputs("arnir0/Tiny-LLM")
-        ds = all_dynamic_shape_from_inputs(data["inputs"])
+        ds = all_dynamic_shapes_from_inputs(data["inputs"])
         self.assertEqual(
             {
                 "input_ids": {0: "d_0_0", 1: "d_0_1"},
@@ -184,6 +185,60 @@ def test_guess_dynamic_shapes_from_inputs(self):
             guessed,
         )
 
+    @requires_transformers("4.55")
+    @requires_torch("2.9")
+    def test_make_fake_with_dynamic_dimensions_tensor(self):
+        res = make_fake_with_dynamic_dimensions(
+            (torch.rand((2, 32, 30, 96), dtype=torch.float16),),
+            ({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])
+
+    @requires_transformers("4.55")
+    @requires_torch("2.9")
+    def test_make_fake_with_dynamic_dimensions_whole(self):
+        res = 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),
+                        ),
+                    ]
+                ),
+            ),
+            dynamic_shapes={
+                "input_ids": {0: "batch", 1: "seq_length"},
+                "attention_mask": {0: "batch", 1: "cache+seq"},
+                "position_ids": {0: "batch", 1: "seq_length"},
+                "past_key_values": [
+                    [{0: "batch", 2: "cache_length"}, {0: "batch", 2: "cache_length"}],
+                    [{0: "batch", 2: "cache_length"}, {0: "batch", 2: "cache_length"}],
+                ],
+            },
+        )
+        flat = flatten_object(res[0], drop_keys=True)
+        for t in flat:
+            if len(t.shape) == 4:
+                self.assertIsInstance(t.shape[0], torch.SymInt)
+                self.assertIsInstance(t.shape[2], torch.SymInt)
+                self.assertEqual(t.shape[1], 32)
+                self.assertEqual(t.shape[3], 96)
+                self.assertNotEqual(t.shape[0], t.shape[2])
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

_unittests/ut_helpers/test_fake_tensor_helper.py

@@ -0,0 +1,74 @@
+import unittest
+import torch
+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
+
+
+class TestMakeTensorHelper(ExtTestCase):
+
+    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"})
+        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"})
+        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"})
+        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"})
+        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)

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.")