From 6eb6adff2f1a2fe988e79c53e622b18f9ff7ab26 Mon Sep 17 00:00:00 2001
From: Xavier Dupre <xadupre@microsoft.com>
Date: Fri, 2 Feb 2024 17:22:54 +0100
Subject: [PATCH 1/2] update requirements

---
 azure-pipelines.yml  | 4 ++--
 requirements-dev.txt | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/azure-pipelines.yml b/azure-pipelines.yml
index 907bb9f..61587f4 100644
--- a/azure-pipelines.yml
+++ b/azure-pipelines.yml
@@ -4,8 +4,8 @@ jobs:
     vmImage: 'ubuntu-latest'
   strategy:
     matrix:
-      Python311-Linux:
-        python.version: '3.11'
+      Python312-Linux:
+        python.version: '3.12'
     maxParallel: 3
 
   steps:
diff --git a/requirements-dev.txt b/requirements-dev.txt
index 5804529..5e262e3 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -11,7 +11,7 @@ lightgbm
 matplotlib
 ml-dtypes
 git+https://github.com/onnx/onnxmltools.git
-onnxruntime>=1.16.1
+onnxruntime>=1.17.0
 openpyxl
 packaging
 pandas

From 17a03aa19c41fd5ee6bc8ead21d4b1b933584a26 Mon Sep 17 00:00:00 2001
From: Xavier Dupre <xadupre@microsoft.com>
Date: Mon, 5 Feb 2024 12:27:43 +0100
Subject: [PATCH 2/2] fix bugs in remove_identity nodes

---
 onnx_array_api/graph_api/graph_builder.py | 41 +++++++++++++++++++----
 pyproject.toml                            |  4 +--
 2 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/onnx_array_api/graph_api/graph_builder.py b/onnx_array_api/graph_api/graph_builder.py
index f238eee..c9c2059 100644
--- a/onnx_array_api/graph_api/graph_builder.py
+++ b/onnx_array_api/graph_api/graph_builder.py
@@ -836,11 +836,12 @@ def remove_identity_nodes(self):
         """
         Removes identity nodes.
         """
-        # f<irst pass: detect replacements
+        # first pass: detect replacements
         new_nodes = []
         input_names = set(i.name for i in self.inputs)
         output_names = set(i.name for i in self.outputs)
         replacements = {}
+        replacements_rev = {}
         for node in self.nodes:
             if node.op_type != "Identity":
                 new_nodes.append(node)
@@ -848,18 +849,44 @@ def remove_identity_nodes(self):
 
             if node.output[0] not in output_names:
                 old_name, new_name = node.output[0], node.input[0]
-            elif node.input[0] not in input_names:
+            elif (
+                node.input[0] not in input_names
+                and node.input[0] not in output_names
+                and node.input[0] not in replacements
+            ):
                 old_name, new_name = node.input[0], node.output[0]
             else:
                 new_nodes.append(node)
                 continue
 
             # the new name can be set for replacements as well
-            assert old_name not in replacements
             if new_name in replacements:
                 new_name = replacements[new_name]
-                assert new_name not in replacements
+                assert new_name not in replacements, (
+                    f"Name {old_name!r} still in {replacements}, node.op_type={node.op_type!r}, "
+                    f"node.input={node.input}, node.output={node.output}, "
+                    f"input_names={input_names}, output_names={output_names}"
+                )
+            if old_name in replacements_rev:
+                old_old_name = replacements_rev[old_name]
+                replacements[old_old_name] = new_name
+                replacements_rev[new_name] = old_old_name
+            if old_name in replacements:
+                replacements[replacements[old_name]] = new_name
+            assert new_name not in replacements, (
+                f"Name {old_name!r} still in {replacements}, node.op_type={node.op_type!r}, "
+                f"node.input={node.input}, node.output={node.output}, "
+                f"input_names={input_names}, output_names={output_names}"
+            )
             replacements[old_name] = new_name
+            replacements_rev[new_name] = old_name
+
+            # verification
+            for k, v in replacements.items():
+                assert v not in replacements, (
+                    f"replacement {k}->{v} is not possible because of "
+                    f"{v}->{replacements[v]}, old_name={old_name!r}, new_name={new_name!r}"
+                )
 
         # second pass: replacements in initializer
         for k, v in replacements.items():
@@ -876,10 +903,12 @@ def remove_identity_nodes(self):
             repo = {o for o in node.output if o in replacements}
             repi = {o for o in node.input if o in replacements}
             if repi or repo:
+                new_inputs = [replacements.get(i, i) for i in node.input]
+                new_outputs = [replacements.get(i, i) for i in node.output]
                 new_node = oh.make_node(
                     node.op_type,
-                    [replacements.get(i, i) for i in node.input],
-                    [replacements.get(i, i) for i in node.output],
+                    new_inputs,
+                    new_outputs,
                     domain=node.domain,
                     name=node.name,
                 )
diff --git a/pyproject.toml b/pyproject.toml
index 0b0e71d..525b648 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -11,11 +11,11 @@ exclude = [
 # Same as Black.
 line-length = 88
 
-[tool.ruff.mccabe]
+[tool.ruff.lint.mccabe]
 # Unlike Flake8, default to a complexity level of 10.
 max-complexity = 10
 
-[tool.ruff.per-file-ignores]
+[tool.ruff.lint.per-file-ignores]
 "_doc/examples/plot_first_example.py" = ["E402", "F811"]
 "_doc/examples/plot_onnxruntime.py" = ["E402", "F811"]
 "onnx_array_api/array_api/_onnx_common.py" = ["F821"]