more about fake helper

Author: xadupre

CHANGELOGS.rst

@@ -4,7 +4,7 @@ Change Logs
 0.7.16
 ++++++
 
-* :pr:`272`: makes patches woth with FakeTensor
+* :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``)
@@ -101,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/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,18 +1,18 @@
 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.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,
+)
 
 
 class TestShapeHelper(ExtTestCase):
@@ -21,7 +21,7 @@ class TestShapeHelper(ExtTestCase):
     @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 +122,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 +145,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"},

_unittests/ut_helpers/test_fake_tensor_helper.py

@@ -3,10 +3,44 @@
 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
+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(

onnx_diagnostic/export/dynamic_shapes.py

@@ -2,24 +2,23 @@
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 import numpy as np
 import torch
-from ..helpers import string_type, flatten_object
+from ..helpers import string_type
 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]]
 
 
-def flatten_dynamic_shapes(ds: Any) -> Any:
+def _flatten_dynamic_shapes(ds: Any) -> Any:
     """Flattens the dynamic shapes."""
     if isinstance(ds, list):
-        return _flat_list([flatten_dynamic_shapes(t) for t in ds])
+        return _flat_list([_flatten_dynamic_shapes(t) for t in ds])
     if isinstance(ds, tuple):
-        return tuple(_flat_list([flatten_dynamic_shapes(t) for t in ds]))
+        return tuple(_flat_list([_flatten_dynamic_shapes(t) for t in ds]))
     if isinstance(ds, dict):
         if all(isinstance(i, int) for i in ds):
             # That's a dynamic shape
             return ds
-        return _flat_list([flatten_dynamic_shapes(t) for t in ds.values()])
+        return _flat_list([_flatten_dynamic_shapes(t) for t in ds.values()])
     raise AssertionError(f"Not implemented for {type(ds)}: {ds}")
 
 
@@ -33,51 +32,6 @@ 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.
@@ -426,7 +380,7 @@ def _generic_walker_step(
         flat, spec = torch.utils._pytree.tree_flatten(inputs)
         if all(isinstance(t, torch.Tensor) for t in flat):
             # We need to flatten dynamic shapes as well
-            ds = flatten_dynamic_shapes(ds)
+            ds = _flatten_dynamic_shapes(ds)
         res = cls._generic_walker_step(
             processor, flat, ds, flatten_unflatten=flatten_unflatten
         )