From 42950912acc37e4c0c8b4c21949ec03a259f4dda Mon Sep 17 00:00:00 2001
From: xadupre <xadupre@microsoft.com>
Date: Fri, 2 May 2025 19:24:20 +0200
Subject: [PATCH 1/8] fix inplace=-1

---
 experimental_experiment/torch_interpreter/export_options.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/experimental_experiment/torch_interpreter/export_options.py b/experimental_experiment/torch_interpreter/export_options.py
index 98158e92..697da1ab 100644
--- a/experimental_experiment/torch_interpreter/export_options.py
+++ b/experimental_experiment/torch_interpreter/export_options.py
@@ -207,7 +207,7 @@ def post_process_exported_program(
                     f"[ExportOptions.export] done remove inplace in "
                     f"{time.perf_counter() - begin}, modified={modified}"
                 )
-            if modified < -1:
+            if modified <= -1:
                 # We need to run decomposition to fully remove all inplace operations.
                 if verbose:
                     begin = time.perf_counter()

From f64036bd19570f3826892b99478c08ff2c43a1a4 Mon Sep 17 00:00:00 2001
From: xadupre <xadupre@microsoft.com>
Date: Fri, 2 May 2025 19:47:52 +0200
Subject: [PATCH 2/8] ci

---
 .github/workflows/documentation.yml |  7 ++++++-
 .github/workflows/wheels-any.yml    |  2 +-
 azure-pipelines.yml                 | 24 ++++++++++++++++++++++++
 3 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/documentation.yml b/.github/workflows/documentation.yml
index 4aefd1d4..03bf2788 100644
--- a/.github/workflows/documentation.yml
+++ b/.github/workflows/documentation.yml
@@ -21,7 +21,7 @@ jobs:
 
       - uses: actions/setup-python@v4
         with:
-          python-version: '3.11'
+          python-version: '3.12'
 
       - uses: tlylt/install-graphviz@v1
 
@@ -62,6 +62,11 @@ jobs:
         run: |
           python -m pip install -r requirements-dev.txt
 
+      - name: Uninstall onnx and install onnx-weekly
+        run: |
+          python -m pip uninstall -y onnx
+          python -m pip install onnx-weekly
+
       - name: Cache pip
         uses: actions/cache@v4
         with:
diff --git a/.github/workflows/wheels-any.yml b/.github/workflows/wheels-any.yml
index 6a99df30..cf03c17a 100644
--- a/.github/workflows/wheels-any.yml
+++ b/.github/workflows/wheels-any.yml
@@ -19,7 +19,7 @@ jobs:
 
       - uses: actions/setup-python@v4
         with:
-          python-version: '3.11'
+          python-version: '3.12'
 
       - name: build wheel
         run: python -m pip wheel .
diff --git a/azure-pipelines.yml b/azure-pipelines.yml
index 7fbbbd39..0868c9b9 100644
--- a/azure-pipelines.yml
+++ b/azure-pipelines.yml
@@ -40,6 +40,10 @@ jobs:
       cd ..
       cd ..
     displayName: 'Install latest transformers'
+  - script: |
+      python -m pip uninstall -y onnx
+      python -m pip install onnx-weekly
+    displayName: 'Uninstall onnx and install onnx-weekly'
   - script: |
       python -c "import torch;print('torch', torch.__version__)"
       python -c "import transformers;print('transformers', transformers.__version__)"
@@ -112,6 +116,10 @@ jobs:
   - script: |
       pip install -r requirements-dev.txt
     displayName: 'Install Requirements dev'
+  - script: |
+      python -m pip uninstall -y onnx
+      python -m pip install onnx-weekly
+    displayName: 'Uninstall onnx and install onnx-weekly'
   - script: |
       python -c "import torch;print('torch', torch.__version__)"
       python -c "import transformers;print('transformers', transformers.__version__)"
@@ -174,6 +182,10 @@ jobs:
     displayName: 'Install Requirements'
   - script: |
       pip install -r requirements-dev.txt
+  - script: |
+      python -m pip uninstall -y onnx
+      python -m pip install onnx-weekly
+    displayName: 'Uninstall onnx and install onnx-weekly'
   - script: |
       pip uninstall -y onnxruntime onnxruntime-training
       pip install onnxruntime-training-cpu
@@ -261,6 +273,10 @@ jobs:
     displayName: 'Install Requirements'
   - script: |
       pip install -r requirements-dev.txt
+  - script: |
+      python -m pip uninstall -y onnx
+      python -m pip install onnx-weekly
+    displayName: 'Uninstall onnx and install onnx-weekly'
   - script: |
       pip uninstall -y onnxruntime onnxruntime-training
       pip install onnxruntime-training-cpu
@@ -351,6 +367,10 @@ jobs:
   - script: |
       pip install -r requirements-dev.txt
     displayName: 'Install Requirements dev'
+  - script: |
+      python -m pip uninstall -y onnx
+      python -m pip install onnx-weekly
+    displayName: 'Uninstall onnx and install onnx-weekly'
   - script: |
       pip install chronos-forecasting --no-deps  # because it enforces transformers to a wrong version
     displayName: 'Install Requirements chronos'
@@ -428,6 +448,10 @@ jobs:
   - script: |
       pip install -r requirements-dev.txt
     displayName: 'Install Requirements dev'
+  - script: |
+      python -m pip uninstall -y onnx
+      python -m pip install onnx-weekly
+    displayName: 'Uninstall onnx and install onnx-weekly'
   - script: pip install onnxmltools --no-deps
     displayName: 'Install onnxmltools'
   - script: |

From 67f16910aa9ea57cc8523c029ac13fdc99b31740 Mon Sep 17 00:00:00 2001
From: xadupre <xadupre@microsoft.com>
Date: Mon, 5 May 2025 16:37:00 +0200
Subject: [PATCH 3/8] fix req

---
 _doc/tutorial/docker.rst | 3 +--
 requirements-dev.txt     | 2 +-
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/_doc/tutorial/docker.rst b/_doc/tutorial/docker.rst
index d0d6bf04..064efbbf 100644
--- a/_doc/tutorial/docker.rst
+++ b/_doc/tutorial/docker.rst
@@ -104,7 +104,6 @@ The docker can be saved with ``sudo docker commit docker-dort-new docker-dort-up
     # optional
     git clone https://github.com/onnx/sklearn-onnx.git
     git clone https://github.com/onnx/onnxmltools.git
-    git clone https://github.com/microsoft/onnxconverter-common.git
 
 **Install ort extension**
 
@@ -136,7 +135,7 @@ Mostly made for research until the ideas migrates to an officially supported pac
 .. code-block:: bash
 
     cd /github/github/experimental-experiment
-    export PYTHONPATH=/github/github/experimental-experiment/:/github/github/onnx-extended:/github/github/onnxscript:/github/github/onnxruntime/build/linux_cuda/Release:/github/github/sklearn-onnx:/github/github/onnxmltools:/github/github/onnxconverter-common
+    export PYTHONPATH=/github/github/experimental-experiment/:/github/github/onnx-extended:/github/github/onnxscript:/github/github/onnxruntime/build/linux_cuda/Release:/github/github/sklearn-onnx:/github/github/onnxmltools
 
     # check that dort is working on llama and export the onnx model (flag --help to see other options)
     python -m experimental_experiment.torch_bench.dort_bench --backend ort+ --device cuda --mixed=1 --export model -w 3 -r 5 --enable_pattern=default+onnxruntime+experimental --num_hidden_layers=1
diff --git a/requirements-dev.txt b/requirements-dev.txt
index edd38415..f0df3e13 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -16,7 +16,7 @@ pytest-coverage
 pytest-subtests
 ruff
 scikit-learn>=1.4.0
-skl2onnx
+git+https://github.com/onnx/sklearn-onnx.git
 sphinx
 sphinx-gallery
 sphinx-issues

From c6eaa230e88cbf69d818dac251a8495d0d70a1c6 Mon Sep 17 00:00:00 2001
From: xadupre <xadupre@microsoft.com>
Date: Mon, 5 May 2025 17:11:02 +0200
Subject: [PATCH 4/8] grad

---
 experimental_experiment/gradient/loss_helper.py           | 8 ++++----
 experimental_experiment/xbuilder/_internal/onnx_export.py | 4 ++--
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/experimental_experiment/gradient/loss_helper.py b/experimental_experiment/gradient/loss_helper.py
index ad962de4..e2a20eca 100644
--- a/experimental_experiment/gradient/loss_helper.py
+++ b/experimental_experiment/gradient/loss_helper.py
@@ -2,7 +2,6 @@
 from typing import Any, Dict, List, Optional, Set, Tuple
 import numpy
 from onnx import ModelProto
-from onnx.mapping import TENSOR_TYPE_TO_NP_TYPE
 from onnx.numpy_helper import to_array
 from onnx.helper import (
     make_node,
@@ -12,6 +11,7 @@
     set_model_props,
 )
 from onnx import TensorProto
+from onnx.helper import tensor_dtype_to_np_dtype
 from ..helpers import from_array_extended
 
 
@@ -113,7 +113,7 @@ def _loss_elastic(
     """
     l1_name = _unique_name(existing_names, "l1_name")
     l2_name = _unique_name(existing_names, "l2_name")
-    dtype = TENSOR_TYPE_TO_NP_TYPE[elem]
+    dtype = tensor_dtype_to_np_dtype(elem)
     onx_l1_weight = from_array_extended(numpy.array([l1_weight], dtype=dtype), name=l1_name)
     onx_l2_weight = from_array_extended(numpy.array([l2_weight], dtype=dtype), name=l2_name)
     inits = [onx_l1_weight, onx_l2_weight]
@@ -168,7 +168,7 @@ def _loss_log(
         raise RuntimeError(  # pragma: no cover
             f"output_name={output_name!r}, log loss does not work on labels."
         )
-    dtype = TENSOR_TYPE_TO_NP_TYPE[elem]
+    dtype = tensor_dtype_to_np_dtype(elem)
     one_name = _unique_name(existing_names, "one_name")
     eps_name = _unique_name(existing_names, "eps_name")
     eps1_name = _unique_name(existing_names, "eps1_name")
@@ -557,7 +557,7 @@ def _replace(ens):
                 v = {"l2": v}
             inits_to_add, nodes_to_add = penalty_loss_onnx(
                 k,
-                dtype=TENSOR_TYPE_TO_NP_TYPE[elem],
+                dtype=tensor_dtype_to_np_dtype(elem),
                 existing_names=existing_names,
                 **v,
             )
diff --git a/experimental_experiment/xbuilder/_internal/onnx_export.py b/experimental_experiment/xbuilder/_internal/onnx_export.py
index 67146f25..b72cc51a 100644
--- a/experimental_experiment/xbuilder/_internal/onnx_export.py
+++ b/experimental_experiment/xbuilder/_internal/onnx_export.py
@@ -4,7 +4,7 @@
 import onnx
 from onnx.helper import printable_graph, make_node, np_dtype_to_tensor_dtype
 from onnx import numpy_helper, ModelProto
-from onnx.mapping import TENSOR_TYPE_TO_NP_TYPE
+from onnx.helper import tensor_dtype_to_np_dtype
 from .onnx_export_templates import get_numpy_template
 from .numpy_helper import make_numpy_code
 
@@ -611,7 +611,7 @@ def rename_name(name, out):
         map=map,
         select_attribute=select_attribute,
         repr=repr,
-        TENSOR_TYPE_TO_NP_TYPE=TENSOR_TYPE_TO_NP_TYPE,
+        tensor_dtype_to_np_dtype=tensor_dtype_to_np_dtype,
         make_numpy_code=lambda *args, **kwargs: make_numpy_code(
             *args, context=context, used=used, mark_inits=mark_inits, **kwargs
         ),

From 59e6eda723132e28f368f361042146d3a295426a Mon Sep 17 00:00:00 2001
From: xadupre <xadupre@microsoft.com>
Date: Mon, 5 May 2025 17:34:50 +0200
Subject: [PATCH 5/8] fix ci

---
 ...raph_pattern_optimization_ort_attention.py |  7 ++++++-
 azure-pipelines.yml                           | 20 +++++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/_unittests/ut_xoptim/test_graph_pattern_optimization_ort_attention.py b/_unittests/ut_xoptim/test_graph_pattern_optimization_ort_attention.py
index 1024809b..46b00ff6 100644
--- a/_unittests/ut_xoptim/test_graph_pattern_optimization_ort_attention.py
+++ b/_unittests/ut_xoptim/test_graph_pattern_optimization_ort_attention.py
@@ -8,7 +8,11 @@
 )
 from onnx_array_api.translate_api.make_helper import make_node_extended
 from experimental_experiment.reference import ExtendedReferenceEvaluator
-from experimental_experiment.ext_test_case import ExtTestCase, requires_cuda
+from experimental_experiment.ext_test_case import (
+    ExtTestCase,
+    requires_cuda,
+    requires_onnxruntime,
+)
 from experimental_experiment.xbuilder.graph_builder import (
     GraphBuilder,
     OptimizationOptions,
@@ -253,6 +257,7 @@ def test_attention_pattern_1_4d_cuda(self):
         self.assertEqualArray(expected[0].ravel(), got[0].ravel(), atol=0.1)
         self.assertEqualArray(expected[0], got[0], atol=0.1)
 
+    @requires_onnxruntime("1.22")
     def test_attention_pattern_1_4d_cpu(self):
         model = self._get_model_attention_1()
         self.dump_onnx("test_attention_pattern_1.noopt.onnx", model)
diff --git a/azure-pipelines.yml b/azure-pipelines.yml
index 0868c9b9..95d8a6c7 100644
--- a/azure-pipelines.yml
+++ b/azure-pipelines.yml
@@ -147,6 +147,10 @@ jobs:
   - script: |
       python -m pip install . -v -v -v
     displayName: 'install wheel'
+  - script: |
+      python -m pip uninstall -y onnx
+      python -m pip install onnx-weekly
+    displayName: 'Uninstall onnx and install onnx-weekly'
   - script: |
       python -m pip freeze
     displayName: 'pip freeze'
@@ -217,6 +221,10 @@ jobs:
   - script: |
       python -m pip install . -v -v -v
     displayName: 'install wheel'
+  - script: |
+      python -m pip uninstall -y onnx
+      python -m pip install onnx-weekly
+    displayName: 'Uninstall onnx and install onnx-weekly'
   - script: |
       python -m pip freeze
     displayName: 'pip freeze'
@@ -308,6 +316,10 @@ jobs:
   - script: |
       python -m pip install . -v -v -v
     displayName: 'install wheel'
+  - script: |
+      python -m pip uninstall -y onnx
+      python -m pip install onnx-weekly
+    displayName: 'Uninstall onnx and install onnx-weekly'
   - script: |
       python -m pip freeze
     displayName: 'pip freeze'
@@ -396,6 +408,10 @@ jobs:
   - script: |
       python -m pip install . -v -v -v
     displayName: 'install wheel'
+  - script: |
+      python -m pip uninstall -y onnx
+      python -m pip install onnx-weekly
+    displayName: 'Uninstall onnx and install onnx-weekly'
   - script: |
       python -m pip freeze
     displayName: 'pip freeze'
@@ -454,6 +470,10 @@ jobs:
     displayName: 'Uninstall onnx and install onnx-weekly'
   - script: pip install onnxmltools --no-deps
     displayName: 'Install onnxmltools'
+  - script: |
+      python -m pip uninstall -y onnx
+      python -m pip install onnx-weekly
+    displayName: 'Uninstall onnx and install onnx-weekly'
   - script: |
       python -m pip freeze
     displayName: 'pip freeze'

From 09075d55fc171b4c86eb57ae9e1390dddc53ed57 Mon Sep 17 00:00:00 2001
From: xadupre <xadupre@microsoft.com>
Date: Mon, 5 May 2025 19:20:14 +0200
Subject: [PATCH 6/8] fix ci

---
 .github/workflows/documentation.yml           |  2 +-
 azure-pipelines.yml                           | 22 +++++++++----------
 experimental_experiment/helpers.py            |  2 +-
 .../torch_interpreter/tracing.py              | 11 +++++++++-
 4 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/.github/workflows/documentation.yml b/.github/workflows/documentation.yml
index 03bf2788..080bda26 100644
--- a/.github/workflows/documentation.yml
+++ b/.github/workflows/documentation.yml
@@ -64,7 +64,7 @@ jobs:
 
       - name: Uninstall onnx and install onnx-weekly
         run: |
-          python -m pip uninstall -y onnx
+          python -m pip uninstall -y onnx onnx-weekly
           python -m pip install onnx-weekly
 
       - name: Cache pip
diff --git a/azure-pipelines.yml b/azure-pipelines.yml
index 95d8a6c7..66be57c1 100644
--- a/azure-pipelines.yml
+++ b/azure-pipelines.yml
@@ -41,7 +41,7 @@ jobs:
       cd ..
     displayName: 'Install latest transformers'
   - script: |
-      python -m pip uninstall -y onnx
+      python -m pip uninstall -y onnx onnx-weekly
       python -m pip install onnx-weekly
     displayName: 'Uninstall onnx and install onnx-weekly'
   - script: |
@@ -117,7 +117,7 @@ jobs:
       pip install -r requirements-dev.txt
     displayName: 'Install Requirements dev'
   - script: |
-      python -m pip uninstall -y onnx
+      python -m pip uninstall -y onnx onnx-weekly
       python -m pip install onnx-weekly
     displayName: 'Uninstall onnx and install onnx-weekly'
   - script: |
@@ -148,7 +148,7 @@ jobs:
       python -m pip install . -v -v -v
     displayName: 'install wheel'
   - script: |
-      python -m pip uninstall -y onnx
+      python -m pip uninstall -y onnx onnx-weekly
       python -m pip install onnx-weekly
     displayName: 'Uninstall onnx and install onnx-weekly'
   - script: |
@@ -187,7 +187,7 @@ jobs:
   - script: |
       pip install -r requirements-dev.txt
   - script: |
-      python -m pip uninstall -y onnx
+      python -m pip uninstall -y onnx onnx-weekly
       python -m pip install onnx-weekly
     displayName: 'Uninstall onnx and install onnx-weekly'
   - script: |
@@ -222,7 +222,7 @@ jobs:
       python -m pip install . -v -v -v
     displayName: 'install wheel'
   - script: |
-      python -m pip uninstall -y onnx
+      python -m pip uninstall -y onnx onnx-weekly
       python -m pip install onnx-weekly
     displayName: 'Uninstall onnx and install onnx-weekly'
   - script: |
@@ -282,7 +282,7 @@ jobs:
   - script: |
       pip install -r requirements-dev.txt
   - script: |
-      python -m pip uninstall -y onnx
+      python -m pip uninstall -y onnx onnx-weekly
       python -m pip install onnx-weekly
     displayName: 'Uninstall onnx and install onnx-weekly'
   - script: |
@@ -317,7 +317,7 @@ jobs:
       python -m pip install . -v -v -v
     displayName: 'install wheel'
   - script: |
-      python -m pip uninstall -y onnx
+      python -m pip uninstall -y onnx onnx-weekly
       python -m pip install onnx-weekly
     displayName: 'Uninstall onnx and install onnx-weekly'
   - script: |
@@ -380,7 +380,7 @@ jobs:
       pip install -r requirements-dev.txt
     displayName: 'Install Requirements dev'
   - script: |
-      python -m pip uninstall -y onnx
+      python -m pip uninstall -y onnx onnx-weekly
       python -m pip install onnx-weekly
     displayName: 'Uninstall onnx and install onnx-weekly'
   - script: |
@@ -409,7 +409,7 @@ jobs:
       python -m pip install . -v -v -v
     displayName: 'install wheel'
   - script: |
-      python -m pip uninstall -y onnx
+      python -m pip uninstall -y onnx onnx-weekly
       python -m pip install onnx-weekly
     displayName: 'Uninstall onnx and install onnx-weekly'
   - script: |
@@ -465,13 +465,13 @@ jobs:
       pip install -r requirements-dev.txt
     displayName: 'Install Requirements dev'
   - script: |
-      python -m pip uninstall -y onnx
+      python -m pip uninstall -y onnx onnx-weekly
       python -m pip install onnx-weekly
     displayName: 'Uninstall onnx and install onnx-weekly'
   - script: pip install onnxmltools --no-deps
     displayName: 'Install onnxmltools'
   - script: |
-      python -m pip uninstall -y onnx
+      python -m pip uninstall -y onnx onnx-weekly
       python -m pip install onnx-weekly
     displayName: 'Uninstall onnx and install onnx-weekly'
   - script: |
diff --git a/experimental_experiment/helpers.py b/experimental_experiment/helpers.py
index db6d405b..04ad211a 100644
--- a/experimental_experiment/helpers.py
+++ b/experimental_experiment/helpers.py
@@ -50,7 +50,7 @@ def size_type(dtype: Any) -> int:
 
     if dtype == np.float64 or dtype == np.int64:
         return 8
-    if dtype == np.float32 or dtype == np.float32:
+    if dtype == np.float32 or dtype == np.float32 or dtype == np.int32:
         return 4
     if dtype == np.float16 or dtype == np.int16:
         return 2
diff --git a/experimental_experiment/torch_interpreter/tracing.py b/experimental_experiment/torch_interpreter/tracing.py
index 5a7e4209..fbff3fee 100644
--- a/experimental_experiment/torch_interpreter/tracing.py
+++ b/experimental_experiment/torch_interpreter/tracing.py
@@ -938,7 +938,7 @@ def _macro_new_node_(n, current_remove, set_item_args, inplace_functions):
                     seen_nodes.add(n)
                     current_remove.append(n)
                 elif aten_name[-1] != "_" and "_." not in aten_name:
-                    # This is not inplace modification so all stored
+                    # This is no inplace modification so all stored
                     # slice operator are cleaned.
                     set_item_args = {}
                     current_remove = []
@@ -956,6 +956,15 @@ def _macro_new_node_(n, current_remove, set_item_args, inplace_functions):
                     set_item_args = {}
                     current_remove = []
                     inplace_functions = []
+                elif aten_name == "aten::masked_fill_.Scalar":
+                    # python -m experimental_experiment.torch_bench.bash_bench_huggingface
+                    # --model MBartForConditionalGeneration --device cuda --dtype float16
+                    # --export custom --opt_patterns default --verbose 1 --quiet 0
+                    raise NotImplementedError(
+                        f"Unable to handle target {aten_name!r} (could probably be ignored) "
+                        f"with args={n.args} in\n{''.join(map(_str, pos_users))}\n----\n"
+                        f"{err_graph}"
+                    )
                 else:
                     raise NotImplementedError(
                         f"Unable to handle target {aten_name!r} with args={n.args} "

From 5bcc92db5b62d84ac9c3ef86e7d2780d56399c08 Mon Sep 17 00:00:00 2001
From: xadupre <xadupre@microsoft.com>
Date: Mon, 5 May 2025 23:20:29 +0200
Subject: [PATCH 7/8] fix issues

---
 .../test_graph_small_model_optimization.py          |  1 +
 .../ut_xrun_doc/test_documentation_examples.py      | 13 ++++++++++---
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/_unittests/ut_xoptim/test_graph_small_model_optimization.py b/_unittests/ut_xoptim/test_graph_small_model_optimization.py
index 2dbab0ea..19652b19 100644
--- a/_unittests/ut_xoptim/test_graph_small_model_optimization.py
+++ b/_unittests/ut_xoptim/test_graph_small_model_optimization.py
@@ -36,6 +36,7 @@ def test_remove_unused_nodes_np(self):
                 [onh.from_array(np.array([2], dtype=np.float32), name="init")],
             ),
             opset_imports=[oh.make_opsetid("", 18)],
+            ir_version=10,
         )
         check_model(model)
 
diff --git a/_unittests/ut_xrun_doc/test_documentation_examples.py b/_unittests/ut_xrun_doc/test_documentation_examples.py
index 8c1bd3df..19e95224 100644
--- a/_unittests/ut_xrun_doc/test_documentation_examples.py
+++ b/_unittests/ut_xrun_doc/test_documentation_examples.py
@@ -206,12 +206,19 @@ def add_test_methods(cls):
                 if pv.Version(onnx_array_api.__version__) < pv.Version("0.3.1"):
                     reason = "requires onnx_array_api>=0.3.1"
 
-            if not reason and name in {
-                "plot_torch_sklearn_201.py",
-            }:
+            if not reason and name in {"plot_torch_sklearn_201.py"}:
                 if pv.Version(torch.__version__) < pv.Version("2.9"):
                     reason = "requires torch>=2.9"
 
+            if not reason and name in {"plot_torch_export_201.py"}:
+                try:
+                    import onnx_array_api
+
+                    if pv.Version(onnx_array_api.__version__) < pv.Version("0.3.2"):
+                        reason = "requires onnx-array-api>=3.2"
+                except ImportError:
+                    reason = "missing onnx-array-pi"
+
             if (
                 not reason
                 and not has_onnxruntime_training()

From 4dc20a53db31b860b97cf5b30f796b8ac3ce851c Mon Sep 17 00:00:00 2001
From: xadupre <xadupre@microsoft.com>
Date: Tue, 6 May 2025 01:01:56 +0200
Subject: [PATCH 8/8] fix

---
 _unittests/ut_torch_interpreter/test_tracing.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/_unittests/ut_torch_interpreter/test_tracing.py b/_unittests/ut_torch_interpreter/test_tracing.py
index 8bb3d0a5..9d49fea6 100644
--- a/_unittests/ut_torch_interpreter/test_tracing.py
+++ b/_unittests/ut_torch_interpreter/test_tracing.py
@@ -464,6 +464,10 @@ def forward(self, x):
             got = mod(*inp)
             self.assertEqualArray(expected, got)
 
+    def test_lookup_op(self):
+        op = torch._library.utils.lookup_op("aten::masked_fill.Scalar")
+        self.assertEqual("aten::masked_fill.Scalar", op.name())
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)