Fix: Remove deprecated .path access in Muon optimizer for TF 2.16+ compatibility

Author: utsab345

keras/src/optimizers/muon.py

@@ -132,13 +132,25 @@ 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) < 4:
+        if not 1 < len(variable.shape) < 5:
             return True
-        if self.exclude_embeddings and "embedding" in variable.path.lower():
+
+        # Get variable identifier (use .name in Keras 3+)
+        var_identifier = variable.name
+
+        # Check if embedding layer should be excluded
+        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:
-            if re.search(keyword, variable.path):
-                return True
+            try:
+                if re.search(keyword, var_identifier):
+                    return True
+            except re.error:
+                # Skip invalid regex patterns in exclude_layers
+                continue
+
         return False
 
     def build(self, var_list):
@@ -161,13 +173,13 @@ def build(self, var_list):
 
         for var in var_list:
             if not self._overwrite_variable_with_gradient(var):
-                self.adam_momentums[var.path] = (
+                self.adam_momentums[var.name] = (
                     self.add_variable_from_reference(
                         reference_variable=var, name="momentum"
                     )
                 )
                 if self._should_use_adamw(var):
-                    self.adam_velocities[var.path] = (
+                    self.adam_velocities[var.name] = (
                         self.add_variable_from_reference(
                             reference_variable=var, name="velocity"
                         )
@@ -183,7 +195,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.path]
+        m = self.adam_momentums[variable.name]
         self.assign_add(m, ops.add(gradient, m * (self.momentum - 1)))
         shape = variable.shape
         if self.nesterov:
@@ -210,8 +222,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.path]
-        v = self.adam_velocities[variable.path]
+        m = self.adam_momentums[variable.name]
+        v = self.adam_velocities[variable.name]
 
         alpha = lr * ops.sqrt(1 - adam_beta_2_power) / (1 - adam_beta_1_power)
 

keras/src/optimizers/muon_test.py

@@ -1,4 +1,5 @@
 import numpy as np
+import tensorflow as tf
 
 from keras.src import backend
 from keras.src import ops
@@ -10,58 +11,59 @@
 
 class MuonTest(testing.TestCase):
     def test_config(self):
-        optimizer = Muon(
-            learning_rate=0.5,
-            epsilon=1e-5,
-        )
+        optimizer = Muon(learning_rate=0.5, epsilon=1e-5)
         self.run_class_serialization_test(optimizer)
 
     def test_Newton_Schulz(self):
         optimizer = Muon()
         tensor_input = ops.array([[0.2499, 0.9105], [0.2655, 0.8824]])
-        except_output = ops.array([[-0.4422, 0.6457], [0.7285, 0.2968]])
+        expected_output = ops.array([[-0.4422, 0.6457], [0.7285, 0.2968]])
         output = optimizer.zeropower_via_newtonschulz5(tensor_input, 5)
-        self.assertAllClose(output, except_output, rtol=1e-3, atol=1e-3)
+        self.assertAllClose(output, expected_output, rtol=1e-3, atol=1e-3)
 
     def test_adamw_single_step(self):
         optimizer = Muon()
         grads = ops.array([1.0, 6.0, 7.0, 2.0])
-        vars = backend.Variable([1.0, 2.0, 3.0, 4.0], name="test_vars")
-        optimizer.build([vars])
-        optimizer._adamw_update_step(grads, vars, 0.5)
-        self.assertAllClose(vars, [0.5, 1.5, 2.5, 3.5], rtol=1e-4, atol=1e-4)
+        var = backend.Variable([1.0, 2.0, 3.0, 4.0], name="test_vars")
+        optimizer.build([var])
+        optimizer._adamw_update_step(grads, var, 0.5)
+        self.assertAllClose(var, [0.5, 1.5, 2.5, 3.5], rtol=1e-4, atol=1e-4)
 
-    def test_should_use_adamw(self):
-        vars = backend.Variable([[1.0, 2.0], [3.0, 4.0]])
-        optimizer = Muon(exclude_layers=["var"])
-        self.assertAllClose(
-            True,
-            optimizer._should_use_adamw(vars),
+    def test_should_use_adamw_excluded_layer(self):
+        """Ensure exclude_layers keyword works and no .path is accessed."""
+        optimizer = Muon(exclude_layers=["dense"])
+        dummy_var = backend.Variable(
+            [[1.0, 2.0], [3.0, 4.0]], name="dense_kernel_0"
         )
-        embeding = Embedding(2, 2)
-        embeding.build()
-        self.assertAllClose(
-            True,
-            optimizer._should_use_adamw(embeding.weights[0]),
-        )
-        vars = backend.Variable([[1.0, 2.0], [3.0, 4.0]])
+        result = optimizer._should_use_adamw(dummy_var)
+        self.assertTrue(result)
+
+    def test_should_use_adamw_embedding(self):
+        """Embedding layer should use AdamW when exclude_embeddings=True."""
+        embedding = Embedding(2, 2)
+        embedding.build()
+        optimizer = Muon(exclude_embeddings=True)
+        result = optimizer._should_use_adamw(embedding.weights[0])
+        self.assertTrue(result)
+
+    def test_should_use_adamw_dimension_rule(self):
+        """Variables with dimensions not between 2–4 use AdamW."""
+        v_1d = backend.Variable([1.0, 2.0], name="v1d")
+        v_5d = backend.Variable(np.zeros((2, 2, 2, 2, 2)), name="v5d")
         optimizer = Muon()
-        self.assertAllClose(
-            False,
-            optimizer._should_use_adamw(vars),
-        )
+        self.assertTrue(optimizer._should_use_adamw(v_1d))
+        self.assertTrue(optimizer._should_use_adamw(v_5d))
+
+    def test_should_use_adamw_dense_layer(self):
+        """2D dense layer weights should use Muon (False)."""
         dense = Dense(2)
         dense.build([None, 2])
-        self.assertAllClose(
-            False,
-            optimizer._should_use_adamw(dense.weights[0]),
-        )
+        optimizer = Muon()
+        result = optimizer._should_use_adamw(dense.weights[0])
+        self.assertFalse(result)
 
     def test_muon_single_step(self):
-        optimizer = Muon(
-            learning_rate=0.5,
-            weight_decay=0,
-        )
+        optimizer = Muon(learning_rate=0.5, weight_decay=0)
         grads = ops.array([[1.0, 6.0], [7.0, 2.0]])
         vars = backend.Variable([[1.0, 2.0], [3.0, 4.0]])
         optimizer.build([vars])
@@ -81,3 +83,15 @@ def test_clip_value(self):
         grad = [np.array([100.0, 100.0])]
         clipped_grad = optimizer._clip_gradients(grad)
         self.assertAllClose(clipped_grad[0], [1.0, 1.0])
+
+    def test_no_path_attribute_error(self):
+        """Ensure compatibility with TF 2.16+ where
+        ResourceVariable has no .path."""
+        optimizer = Muon()
+        var = tf.Variable([1.0, 2.0], name="test_var")
+        # Force-run method that caused AttributeError in issue #21793
+        try:
+            result = optimizer._should_use_adamw(var)
+            self.assertIn(result, [True, False])
+        except AttributeError as e:
+            self.fail(f"Unexpected AttributeError: {e}")