Merge branch 'main' of https://github.com/sdpython/onnx-diagnostic into gemma

Author: xadupre

.github/workflows/ci.yml

@@ -13,6 +13,7 @@ jobs:
     name: to-${{ matrix.torch }}-tr-${{ matrix.transformers }}-ci ${{ matrix.os }}-${{ matrix.python }}
     runs-on: ${{ matrix.os }}
     strategy:
+      fail-fast: false
       matrix:
         os: [ubuntu-latest]
         python: ['3.10', '3.11', '3.12', '3.13']

CHANGELOGS.rst

@@ -1,6 +1,11 @@
 Change Logs
 ===========
 
+0.7.13
+++++++
+
+* :pr:`244`: add a patch to bypass the exception raised when the dynamic dimension is in {0,1}
+
 0.7.12
 ++++++
 

_doc/conf.py

@@ -114,6 +114,8 @@ def linkcode_resolve(domain, info):
 nitpicky = True
 # See also scikit-learn/scikit-learn#26761
 nitpick_ignore = [
+    ("py:class", "_DimHint"),
+    ("py:class", "KeyPath"),
     ("py:class", "ast.Node"),
     ("py:class", "dtype"),
     ("py:class", "False"),

_doc/index.rst

@@ -239,6 +239,7 @@ The function replaces dynamic dimensions defined as strings by
 Older versions
 ==============
 
+* `0.7.13 <../v0.7.13/index.html>`_
 * `0.7.12 <../v0.7.12/index.html>`_
 * `0.7.11 <../v0.7.11/index.html>`_
 * `0.6.3 <../v0.6.3/index.html>`_

_unittests/ut_helpers/test_torch_helper.py

@@ -273,6 +273,7 @@ def test_replace_string_by_dynamic(self):
                 "_DimHint(type=<_DimHintType.DYNAMIC:3>,min=None,max=None,_factory=True)",
                 "DYN",
             )
+            .replace("DimHint(DYNAMIC)", "DYN")
         )
         self.assertEqual(
             "{'input_ids':{0:DYN,1:DYN},'attention_mask':({0:DYN,1:DYN},),'position_ids':[{0:DYN,1:DYN}]}",

_unittests/ut_tasks/test_tasks_image_to_video.py

@@ -17,8 +17,8 @@ class TestTasksImageToVideo(ExtTestCase):
     @hide_stdout()
     @requires_diffusers("0.35")
     @requires_transformers("4.55")
-    @requires_torch("2.8.99")
-    def test_image_to_video(self):
+    @requires_torch("2.10.99")
+    def test_image_to_video_oblivious(self):
         kwargs = {
             "_diffusers_version": "0.34.0.dev0",
             "_class_name": "CosmosTransformer3DModel",
@@ -63,6 +63,53 @@ def test_image_to_video(self):
                 model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
             )
 
+    @hide_stdout()
+    @requires_diffusers("0.35")
+    @requires_transformers("4.55")
+    @requires_torch("2.8.99")
+    def test_image_to_video_not_oblivious(self):
+        kwargs = {
+            "_diffusers_version": "0.34.0.dev0",
+            "_class_name": "CosmosTransformer3DModel",
+            "max_size": [128, 240, 240],
+            "text_embed_dim": 128,
+            "use_cache": True,
+            "in_channels": 3,
+            "out_channels": 16,
+            "num_layers": 2,
+            "model_type": "dia",
+            "patch_size": [1, 2, 2],
+            "rope_scale": [1.0, 3.0, 3.0],
+            "attention_head_dim": 16,
+            "mlp_ratio": 0.4,
+            "initializer_range": 0.02,
+            "num_attention_heads": 16,
+            "is_encoder_decoder": True,
+            "adaln_lora_dim": 16,
+            "concat_padding_mask": True,
+            "extra_pos_embed_type": None,
+        }
+        config = transformers.DiaConfig(**kwargs)
+        mid = "nvidia/Cosmos-Predict2-2B-Video2World"
+        data = get_untrained_model_with_inputs(
+            mid,
+            verbose=1,
+            add_second_input=True,
+            subfolder="transformer",
+            config=config,
+            inputs_kwargs=dict(image_height=8 * 50, image_width=8 * 80),
+        )
+        self.assertEqual(data["task"], "image-to-video")
+        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        model(**inputs)
+        model(**data["inputs2"])
+        with torch_export_patches(
+            patch_transformers=True, patch_diffusers=True, verbose=10, stop_if_static=1
+        ):
+            torch.export.export(
+                model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
+            )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

_unittests/ut_torch_export_patches/test_patch_rewrite.py

@@ -3,6 +3,7 @@
 from onnx_diagnostic.torch_export_patches.patches.patch_transformers import (
     rewrite_loop_for_square_mask,
 )
+from onnx_diagnostic.torch_export_patches.patch_module_helper import code_needing_rewriting
 
 
 class TestPatchRewriting(ExtTestCase):
@@ -33,6 +34,10 @@ def apply_mask(mask, seq):
                 m2 = rewrite_loop_for_square_mask(mask, seq)
                 self.assertEqualArray(m1, m2)
 
+    def test_code_needing_rewriting(self):
+        res = code_needing_rewriting("BartModel")
+        self.assertEqual(len(res), 2)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

_unittests/ut_torch_export_patches/test_patch_rewriting.py

@@ -2,7 +2,14 @@
 from typing import Callable
 import torch
 from torch._dynamo._trace_wrapped_higher_order_op import TransformGetItemToIndex
-from onnx_diagnostic.ext_test_case import ExtTestCase, requires_torch, requires_transformers
+from onnx_diagnostic.ext_test_case import (
+    ExtTestCase,
+    requires_torch,
+    requires_transformers,
+    has_torch,
+)
+from onnx_diagnostic.torch_export_patches import torch_export_patches
+from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
 
 
 class TestPatchPatchTorch(ExtTestCase):
@@ -236,6 +243,79 @@ def forward(self, x):
         ep = torch.export.export(Model(), (x,), dynamic_shapes=({0: DYN},))
         self.assertEqualArray(Model()(x), ep.module()(x))
 
+    def test_oblivious_for_dimension_01(self):
+        class Model(torch.nn.Module):
+            def forward(self, x, ind1, ind2):
+                return x[ind1, ind2]
+
+        inputs = (
+            torch.randn(2, 1024),
+            torch.tensor([[0, 1]], dtype=torch.int64).T,
+            torch.arange(1024, dtype=torch.int64),
+        )
+        model = Model()
+        expected = model(*inputs)
+
+        dynamic_string = ({0: "A", 1: "B"}, {0: "C", 1: "D"}, {0: "E"})
+        # ({0: DYN, 1: DYN}, {0: DYN, 1: DYN}, {0: DYN})
+
+        dynamic_shapes = use_dyn_not_str(dynamic_string)
+        with self.subTest(
+            name="export 0/1 specialized due to hint of 1 for dimension",
+            dynamic_shapes=dynamic_shapes,
+        ):
+            try:
+                torch.export.export(model, inputs, dynamic_shapes=dynamic_shapes)
+                raise AssertionError("torch fixed that case")
+            except ValueError as e:
+                self.assertIn("export 0/1 specialized due to hint of 1 for dimension", str(e))
+
+        dynamic_shapes = use_dyn_not_str(dynamic_string, torch.export.Dim.AUTO)
+        if has_torch("2.9"):
+            with self.subTest(
+                name="expected shape should be broadcastable to (>= 2.9)",
+                dynamic_shapes=dynamic_shapes,
+            ):
+                try:
+                    with torch.fx.experimental._config.patch(backed_size_oblivious=True):
+                        torch.export.export(model, inputs, dynamic_shapes=dynamic_shapes)
+                    raise AssertionError("torch fixed that case")
+                except RuntimeError as e:
+                    self.assertIn("expected shape should be broadcastable to", str(e))
+
+        if not has_torch("2.9"):
+            with self.subTest(
+                name="expected shape should be broadcastable to (< 2.9)",
+                dynamic_shapes=dynamic_shapes,
+            ):
+                try:
+                    with torch.fx.experimental._config.patch(backed_size_oblivious=True):
+                        torch.export.export(model, inputs, dynamic_shapes=dynamic_shapes)
+                except RuntimeError as e:
+                    self.assertIn(
+                        "Expected input at *args[2].shape[0] to be equal to 1, but got 1024",
+                        str(e),
+                    )
+
+        with self.subTest(name="patch for 0/1", dynamic_shapes=dynamic_shapes):
+            with torch_export_patches():
+                ep = torch.export.export(model, inputs, dynamic_shapes=dynamic_shapes)
+            got = ep.module()(*inputs)
+            self.assertEqualArray(expected, got)
+
+        if has_torch("2.11"):
+            # Missing PR https://github.com/pytorch/pytorch/pull/164225
+            # Needs more thinking about the patch to apply for this particular example.
+            with self.subTest(
+                name="patch for 0/1 with oblivious", dynamic_shapes=dynamic_shapes
+            ):
+                with torch_export_patches(), torch.fx.experimental._config.patch(
+                    backed_size_oblivious=True
+                ):
+                    ep = torch.export.export(model, inputs, dynamic_shapes=dynamic_shapes)
+                got = ep.module()(*inputs)
+                self.assertEqualArray(expected, got)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

_unittests/ut_torch_export_patches/test_patch_torch.py

@@ -28,6 +28,7 @@
 
 class TestHuggingFaceHubApi(ExtTestCase):
 
+    @unittest.skip("https://github.com/sdpython/onnx-diagnostic/issues/242")
     @requires_transformers("4.50")  # we limit to some versions of the CI
     @requires_torch("2.7")
     @ignore_errors(OSError)  # connectivity issues