Fix Muon optimizer TF 2.16+ compatibility: replace .path with id(variable) for uniqueness

Author: utsab345

keras/src/optimizers/muon.py

@@ -129,26 +129,27 @@ def __init__(
         self.exclude_layers = exclude_layers or []
 
     def _should_use_adamw(self, variable):
-        # To use it with 4D convolutional filters,
-        # it works well to just flatten their last 3 dimensions.
-        # any {0,1}-D parameters should all be optimized by adam
-        if not 1 < len(variable.shape) < 5:
+        """Determine if AdamW should be used for a variable."""
+        # Use AdamW for variables not having 2D, 3D, or 4D shape
+        if len(variable.shape) not in (2, 3, 4):
             return True
 
-        # Get variable identifier (use .name in Keras 3+)
+        # Get variable identifier (.name for logging and checks)
         var_identifier = variable.name
 
-        # Check if embedding layer should be excluded
+        # Exclude embedding layers
         if self.exclude_embeddings and "embedding" in var_identifier.lower():
             return True
 
-        # Check if variable matches any excluded layer patterns
-        for keyword in self.exclude_layers:
+        # Exclude any user-specified layer patterns
+        for pattern in self.exclude_layers:
             try:
-                if re.search(keyword, var_identifier):
+                if re.search(pattern, var_identifier):
                     return True
-            except re.error:
+            except (re.error, TypeError):
+                # Skip invalid regex patterns or non-string entries
                 continue
+
         return False
 
     def build(self, var_list):
@@ -166,18 +167,13 @@ def build(self, var_list):
         self.adam_momentums = {}
         self.adam_velocities = {}
 
-        self.muon_momentums = {}
-        self.muon_velocities = {}
-
         for var in var_list:
             if not self._overwrite_variable_with_gradient(var):
-                self.adam_momentums[var.name] = (
-                    self.add_variable_from_reference(
-                        reference_variable=var, name="momentum"
-                    )
+                self.adam_momentums[id(var)] = self.add_variable_from_reference(
+                    reference_variable=var, name="momentum"
                 )
                 if self._should_use_adamw(var):
-                    self.adam_velocities[var.name] = (
+                    self.adam_velocities[id(var)] = (
                         self.add_variable_from_reference(
                             reference_variable=var, name="velocity"
                         )
@@ -193,7 +189,7 @@ def update_step(self, gradient, variable, learning_rate):
             self._muon_update_step(gradient, variable, learning_rate)
 
     def _muon_update_step(self, gradient, variable, lr):
-        m = self.adam_momentums[variable.name]
+        m = self.adam_momentums[id(variable)]
         self.assign_add(m, ops.add(gradient, m * (self.momentum - 1)))
         shape = variable.shape
         if self.nesterov:
@@ -220,8 +216,8 @@ def _adamw_update_step(self, gradient, variable, learning_rate):
             ops.cast(self.adam_beta_2, variable.dtype), local_step
         )
 
-        m = self.adam_momentums[variable.name]
-        v = self.adam_velocities[variable.name]
+        m = self.adam_momentums[id(variable)]
+        v = self.adam_velocities[id(variable)]
 
         alpha = lr * ops.sqrt(1 - adam_beta_2_power) / (1 - adam_beta_1_power)
 

keras/src/optimizers/muon_test.py

@@ -96,3 +96,19 @@ def test_no_path_attribute_error(self):
             self.assertIn(result, [True, False])
         except AttributeError as e:
             self.fail(f"Unexpected AttributeError: {e}")
+
+    def test_variable_name_uniqueness(self):
+        """Ensure variable names are unique and no key collisions occur."""
+        optimizer = Muon()
+        # Create variables with different names (simulating real layers)
+        var1 = backend.Variable([[1.0, 2.0]], name="kernel1")
+        var2 = backend.Variable([[3.0, 4.0]], name="kernel2")
+        optimizer.build([var1, var2])
+
+        # Check that each has its own momentum (unique variable objects)
+        self.assertIn(id(var1), optimizer.adam_momentums)
+        self.assertIn(id(var2), optimizer.adam_momentums)
+        self.assertIsNot(
+            optimizer.adam_momentums[id(var1)],
+            optimizer.adam_momentums[id(var2)],
+        )