fix patches

Author: xadupre

_doc/technical/plot_broadcast_export_issue.py

@@ -80,8 +80,8 @@ def forward(self, x, y):
 # 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)
+    torch.zeros((3,), 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)
 res = FakeTensorProp(gm, fake_mode).propagate(*fake_inputs)
@@ -115,7 +115,7 @@ def forward(self, x, y):
 try:
     res = FakeTensorProp(gm, fake_mode).propagate(*fake_inputs)
 except Exception as e:
-    print(e)
+    print("error", e)
 
 # %%
 # By applying the patches:

onnx_diagnostic/_command_lines_parser.py

@@ -693,16 +693,16 @@ def _cmd_export_sample(argv: List[Any]):
             os.makedirs(args.dump_folder, exist_ok=True)
             name = (
                 _make_folder_name(
-                    model_id=args.model_id,
-                    exporter=args.exporter,
-                    optimization=args.optimization,
+                    model_id=args.mid,
+                    exporter=args.export,
+                    optimization=args.opt,
                     dtype=args.dtype,
                     device=args.device,
                     subfolder=args.subfolder,
                     opset=args.opset,
                     drop_inputs=None if not args.drop else args.drop.split(","),
-                    same_as_pretrained=args.same_as_pretrained,
-                    use_pretrained=args.use_pretrained,
+                    same_as_pretrained=args.same_as_trained,
+                    use_pretrained=args.trained,
                     task=args.task,
                 ).replace("/", "-")
                 + ".py"

onnx_diagnostic/torch_export_patches/patches/patch_transformers.py

@@ -1312,6 +1312,10 @@ def patched_sdpa_attention_forward(
     # is_causal = query.shape[2] > 1 and attention_mask is None and is_causal
     is_causal = attention_mask is None and is_causal
 
+    torch._check(
+        attention_mask.shape[3] == key.shape[2],
+        "Attention mask shape incompatible with key shape.",
+    )
     attn_output = torch.nn.functional.scaled_dot_product_attention(
         query,
         key,