fix dynamo

Author: xadupre

CHANGELOGS.rst

@@ -5,6 +5,7 @@ Change Logs
 ++++++
 
 * :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``)
 * :pr:`266`: makes ``patch_torch`` an integer in ``torch_export_patches`` to enable more patches 
 

_doc/technical/plot_broadcast_export_issue.py

@@ -0,0 +1,145 @@
+"""
+Dynamic Shapes and Broadcasting
+===============================
+
+:func:`torch.export.export` makes strict assumption on dynamic shapes
+to the generic case. Let's consider two tensors with only one dimension.
+``x * y`` allows four configurations:
+
+* ``shape(x) = (1,)`` and ``shape(y) = (1,)``
+* ``shape(x) = (1,)`` and ``shape(y) = (p,)``
+* ``shape(x) = (q,)`` and ``shape(y) = (1,)``
+* ``shape(x) = (p,)`` and ``shape(y) = (p,)``
+
+The expected shape for ``shape(x * y)`` is ``(max(p,q),)``.
+
+Simple Case
++++++++++++
+
+"""
+
+import torch
+from torch.fx.experimental.symbolic_shapes import ShapeEnv
+from torch._subclasses.fake_tensor import FakeTensorMode
+from torch.fx.passes.fake_tensor_prop import FakeTensorProp
+from onnx_diagnostic.torch_export_patches import torch_export_patches
+from torch.fx import Tracer
+
+
+class Model(torch.nn.Module):
+    def forward(self, x, y):
+        return x * y
+
+
+Dim = torch.export.Dim
+
+ep = torch.export.export(
+    Model(),
+    (torch.tensor([2, 3], dtype=torch.float32), torch.tensor([2, 3], dtype=torch.float32)),
+    dynamic_shapes=({0: Dim.DYNAMIC}, {0: Dim.DYNAMIC}),
+)
+print(ep)
+
+# %%
+# We see clearly that the export assumed that ``x`` ad ``y`` had the same shape.
+# No other configuration seemed to work at export time,
+# including ``with torch.fx.experimental._config.patch(backed_size_oblivious=True):``
+# the shape of one tensor equal to ``(1,)``.
+
+output = [n for n in ep.graph.nodes if n.op == "output"][0]
+print("output is ", output.name, " arg is", output.args[0])
+
+# %%
+# The final shape is:
+
+shape = output.args[0][0].meta["val"].shape
+print("output shape is ", shape)
+
+# %%
+# Tracing
+# +++++++
+#
+# Let's compare with what a simple tracing would do. Let's use :class:`torch.fx.Tracer`.
+
+graph = Tracer().trace(Model())
+print(graph)
+
+# %%
+output = [n for n in graph.nodes if n.op == "output"][0]
+print("output is ", output.name, " arg is", output.args[0])
+print("The tracer leaves no trace:", output.args[0].__dict__)
+
+# %%
+# Shape propagation
+# +++++++++++++++++
+
+gm = torch.fx.GraphModule(Model(), graph)
+
+shape_env = ShapeEnv()
+fake_mode = FakeTensorMode(shape_env=shape_env)
+# d1 = shape_env.create_unbacked_symint()
+# d2 = shape_env.create_unbacked_symint()
+fake_inputs = fake_mode.from_tensor(
+    torch.zeros((2,), dtype=torch.float32), static_shapes=False
+), fake_mode.from_tensor(torch.zeros((2,), dtype=torch.float32), static_shapes=False)
+
+print("fake_inputs are ", fake_inputs)
+res = FakeTensorProp(gm, fake_mode).propagate(*fake_inputs)
+print("output is", res)
+
+# %%
+# Handle Different Shapes
+# +++++++++++++++++++++++
+
+fake_inputs = fake_mode.from_tensor(
+    torch.zeros((2,), dtype=torch.float32), static_shapes=False
+), fake_mode.from_tensor(torch.zeros((1,), dtype=torch.float32), static_shapes=False)
+
+print("fake_inputs are ", fake_inputs)
+res = FakeTensorProp(gm, fake_mode).propagate(*fake_inputs)
+print("output is", res)
+
+# %%
+# Conclusion
+# ++++++++++
+#
+# We need to give distinct dimensions to get distinct names.
+
+fake_inputs = fake_mode.from_tensor(
+    torch.zeros((2,), dtype=torch.float32), static_shapes=False
+), fake_mode.from_tensor(torch.zeros((3,), dtype=torch.float32), static_shapes=False)
+print("fake_inputs are ", fake_inputs)
+
+
+# %%
+try:
+    res = FakeTensorProp(gm, fake_mode).propagate(*fake_inputs)
+except Exception as e:
+    print(e)
+
+# %%
+# By applying the patches:
+
+with torch_export_patches():
+    res = FakeTensorProp(gm, fake_mode).propagate(*fake_inputs)
+    print("output is", res)
+
+# %%
+# This is what we want. Let's go back to :func:`torch.export.export`
+
+with torch_export_patches():
+    ep = torch.export.export(
+        Model(),
+        (
+            torch.tensor([2, 3], dtype=torch.float32),
+            torch.tensor([2, 3, 4], dtype=torch.float32),
+        ),
+        dynamic_shapes=({0: Dim.DYNAMIC}, {0: Dim.DYNAMIC}),
+    )
+    print(ep)
+
+# %%
+output = [n for n in ep.graph.nodes if n.op == "output"][0]
+print("output is ", output.name, " arg is", output.args[0])
+shape = output.args[0][0].meta["val"].shape
+print("output shape is ", shape)

_unittests/ut_tasks/test_tasks_text_generation.py

@@ -0,0 +1,33 @@
+import unittest
+import torch
+from onnx_diagnostic.ext_test_case import (
+    ExtTestCase,
+    hide_stdout,
+    requires_transformers,
+    requires_torch,
+)
+from onnx_diagnostic.helpers.torch_helper import torch_deepcopy
+from onnx_diagnostic.torch_models.hghub.model_inputs import get_untrained_model_with_inputs
+from onnx_diagnostic.torch_export_patches import torch_export_patches
+from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
+
+
+class TestTasksTextGeneration(ExtTestCase):
+    @hide_stdout()
+    @requires_transformers("4.53")
+    @requires_torch("2.7.99")
+    def test_image_text_to_text_gemma3_for_causallm(self):
+        mid = "hf-internal-testing/tiny-random-Gemma3ForCausalLM"
+        data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True)
+        self.assertEqual(data["task"], "text-generation")
+        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        model(**torch_deepcopy(inputs))
+        model(**data["inputs2"])
+        with torch_export_patches(patch_transformers=True, verbose=10, patch_torch=False):
+            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_torch.py

@@ -491,6 +491,25 @@ def forward(self, x, ind1, ind2):
             )
         self.assertEqualArray(expected, ep.module()(*inputs))
 
+    def test_broadcast_max(self):
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                return x * y
+
+        Dim = torch.export.Dim
+        with torch_export_patches():
+            ep = torch.export.export(
+                Model(),
+                (
+                    torch.tensor([2, 3], dtype=torch.float32),
+                    torch.tensor([2, 3, 4], dtype=torch.float32),
+                ),
+                dynamic_shapes=({0: Dim.DYNAMIC}, {0: Dim.DYNAMIC}),
+            )
+        output = [n for n in ep.graph.nodes if n.op == "output"]
+        shape = output[0].args[0][0].meta["val"].shape
+        self.assertEqual(str(shape), "torch.Size([Max(s17, s77)])")
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/torch_models/code_sample.py

@@ -132,7 +132,6 @@ def make_export_code(
         if opset:
             args.append(f"opset_version={opset}")
         sargs = ", ".join(args)
-        imports = []
         code.extend([f"epo = torch.onnx.export(model, args=(), kwargs=inputs, {sargs})"])
         if optimization:
             imports.append("import onnxscript")

onnx_diagnostic/torch_models/hghub/hub_data_cached_configs.py

@@ -4865,3 +4865,41 @@ def _ccached_google_gemma_3_4b_it_like():
             },
         }
     )
+
+
+def _ccached_hf_internal_testing_tiny_random_gemma3_for_causal_lm():
+    "hf-internal-testing/tiny-random-Gemma3ForCausalLM"
+    return transformers.Gemma3TextConfig(
+        **{
+            "architectures": ["Gemma3ForCausalLM"],
+            "attention_bias": false,
+            "attention_dropout": 0.0,
+            "attn_logit_softcapping": null,
+            "bos_token_id": 2,
+            "cache_implementation": "hybrid",
+            "eos_token_id": [1, 106],
+            "final_logit_softcapping": null,
+            "head_dim": 8,
+            "hidden_activation": "gelu_pytorch_tanh",
+            "hidden_size": 16,
+            "initializer_range": 0.02,
+            "intermediate_size": 32,
+            "max_position_embeddings": 32768,
+            "model_type": "gemma3_text",
+            "num_attention_heads": 2,
+            "num_hidden_layers": 2,
+            "num_key_value_heads": 1,
+            "pad_token_id": 0,
+            "query_pre_attn_scalar": 256,
+            "rms_norm_eps": 1e-06,
+            "rope_local_base_freq": 10000,
+            "rope_scaling": null,
+            "rope_theta": 1000000,
+            "sliding_window": 512,
+            "sliding_window_pattern": 6,
+            "torch_dtype": "float32",
+            "transformers_version": "4.52.0.dev0",
+            "use_cache": true,
+            "vocab_size": 262144,
+        }
+    )

pyproject.toml

@@ -150,7 +150,7 @@ select = [
 "_doc/notebooks/plot_*.py" = ["E402", "B018", "PIE808", "SIM105", "SIM117"]
 "_doc/recipes/plot_*.py" = ["E402", "B018", "PIE808", "SIM105", "SIM117"]
 "_scripts/compare_model_execution.py" = ["E402", "F401"]
-"_doc/technical/plot_*.py" = ["E402", "B018", "PIE808", "SIM105", "SIM117"]
+"_doc/technical/plot_*.py" = ["E402", "B018", "PIE808", "RUF015", "SIM105", "SIM117"]
 "_unittests/*/test*.py" = ["B008", "B904", "PIE808", "SIM117", "SIM105", "UP008"]
 "onnx_diagnostic/export/__init__.py" = ["F401"]
 "onnx_diagnostic/helpers/__init__.py" = ["F401"]