fixes

Author: xadupre

CHANGELOGS.rst

@@ -4,6 +4,8 @@ Change Logs
 0.7.10
 ++++++
 
+* :pr:`218`: patches used sdpa_mask_recent_torch used from _vmap_for_bhqkv
+
 0.7.9
 +++++
 

_unittests/ut_export/test_jit.py

@@ -7,7 +7,6 @@
     ignore_warnings,
     requires_onnxscript,
 )
-from onnx_diagnostic.reference import ExtendedReferenceEvaluator
 from onnx_diagnostic.helpers.torch_helper import is_torchdynamo_exporting
 
 try:
@@ -75,19 +74,6 @@ def forward(self, images, position):
         y = torch.arange(5, dtype=torch.int64) + 1
         expected = model(x, y)
 
-        name = self.get_dump_file("test_export_loop_onnxscript.onnx")
-        torch.onnx.export(
-            model,
-            (x, y),
-            name,
-            dynamic_axes={"images": {0: "batch", 1: "maxdim"}, "position": {0: "batch"}},
-            dynamo=False,
-        )
-        ref = ExtendedReferenceEvaluator(name)
-        feeds = dict(images=x.numpy(), position=y.numpy())
-        got = ref.run(None, feeds)[0]
-        self.assertEqualArray(expected, got)
-
         DYN = torch.export.Dim.DYNAMIC
         ep = torch.export.export(
             model,

_unittests/ut_torch_export_patches/test_dynamic_class.py

@@ -14,6 +14,7 @@
 from onnx_diagnostic.torch_export_patches.onnx_export_errors import (
     torch_export_patches,
 )
+from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
 from onnx_diagnostic.torch_models.hghub.model_inputs import get_untrained_model_with_inputs
 
 
@@ -305,7 +306,7 @@ def test_phi2_export_module(self):
                 model,
                 (),
                 kwargs=inputs,
-                dynamic_shapes=dyn_shapes,
+                dynamic_shapes=use_dyn_not_str(dyn_shapes),
                 strict=False,  # True works but then the it fails during the execution
             )
             # ep = ep.run_decompositions()
@@ -343,7 +344,7 @@ def test_phi2_export_interpreter(self):
                 model,
                 (),
                 kwargs=inputs,
-                dynamic_shapes=dyn_shapes,
+                dynamic_shapes=use_dyn_not_str(dyn_shapes),
                 strict=False,  # True works but then the it fails during the execution
             )
             # ep = ep.run_decompositions()

_unittests/ut_torch_export_patches/test_patch_module.py

@@ -604,8 +604,8 @@ def loop_body_1(z, iv, x, y):
         rewritten_expected2 = RewrittenModel2()(x, y)
         self.assertEqualArray(expected, rewritten_expected2)
 
-        if not has_torch("2.9"):
-            raise unittest.SkipTest("skipped export, torch must be >= 2.9")
+        if not has_torch("2.10"):
+            raise unittest.SkipTest("skipped export, torch must be >= 2.10")
 
         torch.export.export(RewrittenModel2(), (x, y), dynamic_shapes=ds, strict=False)
         ep = torch.export.export(Model(), (x, y), dynamic_shapes=ds, strict=False)

_unittests/ut_torch_export_patches/test_patch_torch.py

@@ -17,7 +17,7 @@ def test_vmap(self):
         got = patched_vmap(f)(x, y)
         self.assertEqualArray(expected, got)
 
-    @requires_torch("2.9")
+    @requires_torch("2.10")
     def test_export_vmap(self):
         class Model(torch.nn.Module):
             def forward(self, x, y):
@@ -206,10 +206,11 @@ def _vmap_for_bhqkv(mask_function: Callable, bh_indices: bool = True) -> Callabl
 
         class Model(torch.nn.Module):
             def forward(self, batch_arange, head_arange, cache_position, kv_arange):
-                with TransformGetItemToIndex():
-                    causal_mask2 = _vmap_for_bhqkv2(mask_function)(
-                        batch_arange, head_arange, cache_position, kv_arange
-                    )
+                # with TransformGetItemToIndex():
+                # This context as ignored in 2.8 and not any more in 2.9.
+                causal_mask2 = _vmap_for_bhqkv2(mask_function)(
+                    batch_arange, head_arange, cache_position, kv_arange
+                )
                 return causal_mask2
 
         inputs = batch_arange, head_arange, cache_position, kv_arange

_unittests/ut_torch_models/test_validate_whole_models.py

@@ -96,7 +96,7 @@ def test_f_validate_model_onnx_dynamo_ir(self):
         )
 
     @requires_torch("2.7")
-    @requires_onnxscript("0.5")
+    @requires_onnxscript("0.7")
     @hide_stdout()
     @ignore_warnings(FutureWarning)
     def test_g_validate_model_onnx_dynamo_os_ort(self):