fix

Author: xadupre

_unittests/ut_torch_export_patches/test_patch_torch.py

@@ -8,6 +8,8 @@
     requires_transformers,
     has_torch,
 )
+from onnx_diagnostic.helpers.torch_helper import torch_deepcopy
+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.torch_export_patches.patch_inputs import use_dyn_not_str
 
@@ -341,6 +343,19 @@ def forward(self, x, ind1, ind2):
             )
         self.assertEqualArray(expected, ep.module()(*inputs), atol=1e-2)
 
+    @requires_torch("2.7.9999")
+    @requires_transformers("4.49.9999")
+    def test_export_tiny_llm_dim_meta(self):
+        data = get_untrained_model_with_inputs("arnir0/Tiny-LLM", verbose=0)
+        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        expected = model(**torch_deepcopy(inputs))
+        with torch_export_patches(patch_transformers=True):
+            ep = torch.export.export(
+                model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds)
+            )
+        got = ep.module()(**inputs)
+        self.assertEqualArrayAny(expected, got)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -386,8 +386,10 @@ def torch_export_patches(
             )
 
             # torch._prims._broadcast_in_dim_meta
+            f_broadcast_in_dim = torch._prims.broadcast_in_dim
             f__broadcast_in_dim_meta = torch._prims._broadcast_in_dim_meta
             torch._prims._broadcast_in_dim_meta = patched__broadcast_in_dim_meta
+            torch._prims.broadcast_in_dim = patched__broadcast_in_dim_meta
 
             # torch._refs._maybe_broadcast
             f__maybe_broadcast = torch._refs._maybe_broadcast
@@ -595,6 +597,7 @@ def torch_export_patches(
                     f___constrain_user_specified_dimhint_range
                 )
                 torch._prims._broadcast_in_dim_meta = f__broadcast_in_dim_meta
+                torch._prims.broadcast_in_dim = f_broadcast_in_dim
                 torch._refs._maybe_broadcast = f__maybe_broadcast
 
                 if verbose:
@@ -735,9 +738,7 @@ def torch_export_patches(
 
 
 def replacement_before_exporting(args: Any) -> Any:
-    """
-    Does replacements on the given inputs if needed.
-    """
+    """Does replacements on the given inputs if needed."""
     if args is None:
         return None
     if isinstance(args, (int, float)):

onnx_diagnostic/torch_export_patches/patches/patch_torch.py

@@ -608,7 +608,7 @@ def should_expand(a: ShapeType, b: ShapeType) -> bool:
 
             # u0==u1 assume the same, no broadcasting!
             # PATCHED: avoid errors
-            return x != y
+            return True  # guard_or_true(x != y)
             # torch._check(
             #    x == y,
             #    lambda x=x, y=y: (
@@ -665,12 +665,13 @@ def patched__broadcast_in_dim_meta(
     # (no relative reordering of dims) of integers and
     # each dimension must be within the new shape
     def _greater_than_reduce(acc, x):
-        assert isinstance(x, torch.export.Dim)
+        assert isinstance(x, (int, torch.export.Dim)), f"unexpected type {type(x)} for x"
         assert x > acc
         assert x < len(shape)
 
         return x
 
+    print("****", broadcast_dimensions)
     reduce(_greater_than_reduce, broadcast_dimensions, -1)
 
     # shape must be broadcastable to
@@ -694,13 +695,14 @@ def _greater_than_reduce(acc, x):
                 else:
                     new_strides.append(0)
             else:
-                torch._check(
-                    a.shape[original_idx] == shape[idx],
-                    lambda idx=idx, original_idx=original_idx: (
-                        f"non-broadcasting semantics require "
-                        f"{a.shape[original_idx]} == {shape[idx]}"
-                    ),
-                )
+                # PATCHED: disabled this check
+                # torch._check(
+                #    a.shape[original_idx] == shape[idx],
+                #    lambda idx=idx, original_idx=original_idx: (
+                #        f"non-broadcasting semantics require "
+                #        f"{a.shape[original_idx]} == {shape[idx]}"
+                #    ),
+                # )
                 new_strides.append(a.stride()[original_idx])
             original_idx = original_idx + 1
         else: