Address PR review feedback on rename_values

Signed-off-by: enpasos <matthias.unverzagt@enpasos.com>

Author: enpasos

src/onnx_ir/_convenience/__init__.py

@@ -363,63 +363,61 @@ def replace_all_uses_with(
 
 def rename_values(
     values: _core.Value | Sequence[_core.Value],
-    names: str | None | Sequence[str | None],
+    names: str | Sequence[str],
 ) -> None:
     """Rename one or more values.
 
-    When multiple initializer-backed values are renamed together, temporary names are
-    used so swaps and other permutations do not trip the initializer name guard.
+    Initializer-backed values are removed from their graphs while renaming so swaps
+    and other permutations do not trip the initializer name guard.
 
     Args:
         values: The value or values to rename. Must be `_core.Value` instances or a
             sequence of `_core.Value` instances.
         names: The target name or names.
 
     Raises:
+        TypeError: If a value is not a :class:`~onnx_ir.Value` or a name is not a string.
         ValueError: If the number of values and names do not match, if one value is
             given conflicting target names, or if an initializer target would collide
             with an initializer outside the renamed set.
     """
     if not isinstance(values, Sequence):
         values = (values,)
-    if isinstance(names, str) or names is None:
+    if isinstance(names, str) or not isinstance(names, Sequence):
         names = (names,)
     if len(values) != len(names):
         raise ValueError("The number of values and names must match.")
 
-    ordered_pairs: list[tuple[_core.Value, str | None]] = []
-    target_by_value_id: dict[int, str | None] = {}
+    ordered_pairs: list[tuple[_core.Value, str]] = []
+    target_by_value: dict[_core.Value, str] = {}
     for value, name in zip(values, names):
         if not isinstance(value, _core.Value):
             raise TypeError(f"value must be a Value object, not {type(value)}")
-        value_id = id(value)
-        if value_id in target_by_value_id:
-            if target_by_value_id[value_id] != name:
+        if not isinstance(name, str):
+            raise TypeError(f"name must be a string, not {type(name)}")
+        if value in target_by_value:
+            if target_by_value[value] != name:
                 raise ValueError(
                     f"Conflicting target names for value {value!r}: "
-                    f"{target_by_value_id[value_id]!r} vs {name!r}."
+                    f"{target_by_value[value]!r} vs {name!r}."
                 )
             continue
-        target_by_value_id[value_id] = name
+        target_by_value[value] = name
         ordered_pairs.append((value, name))
 
-    initializer_pairs_by_graph: dict[_core.Graph, list[tuple[_core.Value, str | None]]] = {}
+    initializer_pairs_by_graph: dict[_core.Graph, list[tuple[_core.Value, str]]] = {}
     for value, name in ordered_pairs:
         if not value.is_initializer():
             continue
         graph = value.graph
         assert isinstance(graph, _core.Graph), "Initializer values must belong to a graph"
         initializer_pairs_by_graph.setdefault(graph, []).append((value, name))
 
+    initializer_values_by_graph: dict[_core.Graph, tuple[_core.Value, ...]] = {}
     for graph, initializer_pairs in initializer_pairs_by_graph.items():
-        renamed_initializer_ids = {id(value) for value, _ in initializer_pairs}
+        renamed_initializers = {value for value, _ in initializer_pairs}
         seen_targets: dict[str, _core.Value] = {}
         for value, name in initializer_pairs:
-            if name is None:
-                raise ValueError(
-                    "Initializer value cannot have name set to None. "
-                    "Please pop() the value from initializers first to do so."
-                )
             existing = seen_targets.get(name)
             if existing is not None and existing is not value:
                 raise ValueError(
@@ -429,34 +427,26 @@ def rename_values(
             seen_targets[name] = value
             if name in graph.initializers:
                 existing_initializer = graph.initializers[name]
-                if (
-                    existing_initializer is not value
-                    and id(existing_initializer) not in renamed_initializer_ids
-                ):
+                if existing_initializer is not value and existing_initializer not in renamed_initializers:
                     raise ValueError(
                         f"Cannot rename initializer '{value}' to '{name}': "
                         "an initializer with that name already exists."
                     )
 
-        used_names = set(graph.initializers)
-        used_names.update(
-            target_name for _, target_name in initializer_pairs if target_name is not None
-        )
-        tmp_index = 0
-        for value, name in initializer_pairs:
-            if value.name == name:
-                continue
-            tmp_name = f"__onnx_ir_tmp_name_{tmp_index}"
-            while tmp_name in used_names:
-                tmp_index += 1
-                tmp_name = f"__onnx_ir_tmp_name_{tmp_index}"
-            tmp_index += 1
-            used_names.add(tmp_name)
-            value.name = tmp_name
+        initializer_values_by_graph[graph] = tuple(value for value, _ in initializer_pairs)
+
+    for graph, initializer_values in initializer_values_by_graph.items():
+        for value in initializer_values:
+            assert value.name is not None, "Initializer values must have names"
+            graph.initializers.pop(value.name)
 
     for value, name in ordered_pairs:
         value.name = name
 
+    for graph, initializer_values in initializer_values_by_graph.items():
+        for value in initializer_values:
+            graph.initializers.add(value)
+
 
 def create_value_mapping(
     graph: _core.Graph | _core.GraphView | _core.Function,

src/onnx_ir/_convenience/_convenience_test.py

@@ -121,6 +121,29 @@ def test_rename_values_supports_initializer_swaps(self):
         self.assertIs(graph.initializers["const_1"], first)
         self.assertIs(graph.initializers["const_0"], second)
 
+    def test_rename_values_rejects_none_names(self):
+        value = ir.Value(name="value")
+
+        with self.assertRaisesRegex(TypeError, "name must be a string"):
+            _convenience.rename_values(value, None)
+
+    def test_rename_values_rejects_initializer_collisions_outside_rename_set(self):
+        first = ir.Value(name="const_0", const_value=ir.tensor([1], name="const_0"))
+        second = ir.Value(name="const_1", const_value=ir.tensor([2], name="const_1"))
+        graph = ir.Graph(
+            inputs=(),
+            outputs=[first, second],
+            nodes=(),
+            initializers=[first, second],
+            name="test_graph",
+        )
+
+        with self.assertRaisesRegex(ValueError, "an initializer with that name already exists"):
+            _convenience.rename_values(first, "const_1")
+
+        self.assertIs(graph.initializers["const_0"], first)
+        self.assertIs(graph.initializers["const_1"], second)
+
 
 if __name__ == "__main__":
     unittest.main()