Adds unit test and example for broadcast rules related to dimension (#269)

* add unit test and example for boradcast_max

* documentation

Author: xadupre

CHANGELOGS.rst

@@ -4,6 +4,7 @@ Change Logs
 0.7.16
 ++++++
 
+* :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_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)

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"]