Update the graph initializers dictionary when renaming Value (#197)

This pull request enhances the handling of value renaming within the
ONNX IR core: When a value (especially an initializer) is renamed, all
relevant references—including the graph's initializer dictionary—are
consistently updated. This reduces the risk of stale or inconsistent
state when manipulating graph initializers.

Key changes:

**Core logic improvements:**

* The `Value` class now ensures that renaming a value also updates the
corresponding entry in the graph's `initializers` dictionary,
maintaining consistency between the value's name and its registration in
the graph. An error is raised if an initializer's name is set to `None`.

* The `__init__` method docstring for `Value` has been updated to
document this new behavior, clarifying that renaming a value also
updates the graph initializer entry if applicable.

**Code simplification:**

* In `naming.py`, redundant logic for manually updating the
`initializers` dictionary after a value rename has been removed, since
this is now handled automatically by the `Value` class. This makes the
code cleaner and less error-prone.

---------

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: justinchuby

src/onnx_ir/_core.py

@@ -2031,6 +2031,14 @@ def __init__(
     ) -> None:
         """Initialize a value.
 
+        When assigning a name to the value, the name of the backing `const_value` (Tensor)
+        will also be updated. If the value is an initializer of a graph, the initializers
+        dictionary of the graph will also be updated.
+
+        .. versionchanged:: 0.1.10
+            Assigning a name to the value will also update the graph initializer entry
+            if the value is an initializer of a graph.
+
         Args:
             producer: The node that produces the value.
                 It can be ``None`` when the value is initialized first than its producer.
@@ -2179,7 +2187,23 @@ def name(self) -> str | None:
     def name(self, value: str | None) -> None:
         if self._const_value is not None:
             self._const_value.name = value
+        old_name = self._name
         self._name = value
+        if self.is_initializer():
+            if value is None:
+                raise ValueError(
+                    "Initializer value cannot have name set to None. Please pop() the value from initializers first"
+                )
+            # Rename the initializer entry in the graph
+            graph = self._graph
+            assert graph is not None
+            assert old_name is not None
+            if value in graph.initializers and graph.initializers[value] is not self:
+                raise ValueError(
+                    f"Cannot rename initializer to '{value}': an initializer with that name already exists."
+                )
+            graph.initializers.pop(old_name)
+            graph.initializers[value] = self
 
     @property
     def type(self) -> _protocols.TypeProtocol | None:

src/onnx_ir/_core_test.py

@@ -779,6 +779,108 @@ def test_consumers(self):
         self.assertEqual(self.node.outputs[0].consumers(), ())
         self.assertEqual(self.node.outputs[1].consumers(), ())
 
+    def test_name_setter_updates_const_value_name(self):
+        """Test that setting a Value's name also updates the const_value's name if it exists."""
+        tensor = ir.tensor([1, 2, 3], name="original_tensor_name")
+        value = _core.Value(name="original_value_name", const_value=tensor)
+
+        # Verify initial state
+        self.assertEqual(value.name, "original_value_name")
+        self.assertEqual(value.const_value.name, "original_tensor_name")
+
+        # Update the value's name and verify const_value name is also updated
+        value.name = "new_name"
+        self.assertEqual(value.name, "new_name")
+        self.assertEqual(value.const_value.name, "new_name")
+
+        # Test setting name to None
+        value.name = None
+        self.assertIsNone(value.name)
+        self.assertIsNone(value.const_value.name)
+
+    def test_name_setter_without_const_value(self):
+        """Test that setting a Value's name works normally when no const_value exists."""
+        value = _core.Value(name="original_name")
+
+        # Verify initial state
+        self.assertEqual(value.name, "original_name")
+        self.assertIsNone(value.const_value)
+
+        # Update the name
+        value.name = "new_name"
+        self.assertEqual(value.name, "new_name")
+
+        # Set to None
+        value.name = None
+        self.assertIsNone(value.name)
+
+    def test_initializer_name_setter_raises_when_set_to_none(self):
+        """Test that setting an initializer value's name to None raises ValueError."""
+        tensor = ir.tensor([1, 2, 3])
+        value = _core.Value(name="initializer1", const_value=tensor)
+        _core.Graph(inputs=(), outputs=(), nodes=(), initializers=[value])
+
+        # Verify the value is an initializer
+        self.assertTrue(value.is_initializer())
+
+        # Attempt to set name to None should raise ValueError
+        with self.assertRaisesRegex(
+            ValueError,
+            "Initializer value cannot have name set to None. Please pop\\(\\) the value from initializers first",
+        ):
+            value.name = None
+
+    def test_initializer_name_setter_updates_graph_initializers_dict(self):
+        """Test that renaming an initializer value updates the graph's initializers dictionary."""
+        tensor = ir.tensor([1, 2, 3])
+        value = _core.Value(name="old_name", const_value=tensor)
+        graph = _core.Graph(inputs=(), outputs=(), nodes=(), initializers=[value])
+
+        # Verify initial state
+        self.assertTrue(value.is_initializer())
+        self.assertIn("old_name", graph.initializers)
+        self.assertIs(graph.initializers["old_name"], value)
+        self.assertEqual(value.name, "old_name")
+
+        # Rename the value and verify the graph's initializers dict is updated
+        value.name = "new_name"
+
+        # Old key should be removed, new key should be added
+        self.assertNotIn("old_name", graph.initializers)
+        self.assertIn("new_name", graph.initializers)
+        self.assertIs(graph.initializers["new_name"], value)
+        self.assertEqual(value.name, "new_name")
+        self.assertEqual(value.const_value.name, "new_name")
+
+    def test_non_initializer_name_setter_works_normally(self):
+        """Test that name changes work normally for values that are not initializers."""
+        # Test regular value (not part of any graph)
+        tensor = ir.tensor([1, 2, 3])
+        value = _core.Value(name="original_name", const_value=tensor)
+
+        self.assertFalse(value.is_initializer())
+
+        # Should be able to change name without issues
+        value.name = "new_name"
+        self.assertEqual(value.name, "new_name")
+        self.assertEqual(value.const_value.name, "new_name")
+
+        # Should be able to set to None without issues
+        value.name = None
+        self.assertIsNone(value.name)
+        self.assertIsNone(value.const_value.name)
+
+        # Test graph input
+        input_value = _core.Value(name="input1")
+        _core.Graph(inputs=[input_value], outputs=(), nodes=())
+
+        self.assertTrue(input_value.is_graph_input())
+        self.assertFalse(input_value.is_initializer())
+
+        # Should be able to rename input without issues
+        input_value.name = "renamed_input"
+        self.assertEqual(input_value.name, "renamed_input")
+
     # TODO(justinchuby): Test all methods
 
 

src/onnx_ir/passes/common/naming.py

@@ -193,14 +193,8 @@ def _process_value(
         if not value.name:
             modified = self._assign_value_name(value, used_value_names, value_counter)
         else:
-            old_name = value.name
             modified = self._fix_duplicate_value_name(value, used_value_names, value_counter)
-            if modified:
-                assert value.graph is not None
-                if value.is_initializer():
-                    value.graph.initializers.pop(old_name)
-                    # Add the initializer back with the new name
-                    value.graph.initializers.add(value)
+            # initializers dictionary is updated automatically when the Value is renamed
 
         # Record the final name for this value
         assert value.name is not None