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

Author: utsab345

keras/src/optimizers/muon.py

@@ -128,19 +128,35 @@ def __init__(
         self.exclude_embeddings = exclude_embeddings
         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) < 4:
+        """
+        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 
+        """
+        # Use Adam for scalar or vector parameters
+        if not 1 < len(variable.shape) <5:
             return True
-        if self.exclude_embeddings and "embedding" in variable.path.lower():
+
+        # Exclude embedding layers if specified
+        var_identifier = getattr(variable, "name", "") or getattr(variable, "path", "")
+        if self.exclude_embeddings and "embedding" in var_identifier.lower():
             return True
-        for keyword in self.exclude_layers:
-            if re.search(keyword, variable.path):
-                return True
+
+        # Exclude variables matching any of the excluded layer patterns
+        for keyword in getattr(self, "exclude_layers", []):
+            try:
+                if re.search(keyword, var_identifier):
+                    return True
+            except re.error:
+                # Skip invalid regex patterns
+                continue
+
+        # Otherwise, use AdamW
         return False
 
+
     def build(self, var_list):
         """Initialize optimizer variables.
 
@@ -161,13 +177,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 +199,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 +226,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

@@ -3,72 +3,65 @@
 from keras.src import backend
 from keras.src import ops
 from keras.src import testing
-from keras.src.layers import Dense
-from keras.src.layers import Embedding
+from keras.src.layers import Dense, Embedding
 from keras.src.optimizers.muon import Muon
 
 
 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]])
+        # Excluded layer test
+        var = backend.Variable([[1.0, 2.0], [3.0, 4.0]])
         optimizer = Muon(exclude_layers=["var"])
-        self.assertAllClose(
-            True,
-            optimizer._should_use_adamw(vars),
-        )
-        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]])
+        self.assertTrue(optimizer._should_use_adamw(var))
+
+        # Embedding test
+        embedding = Embedding(2, 2)
+        embedding.build()
+        optimizer = Muon(exclude_embeddings=True)
+        self.assertTrue(optimizer._should_use_adamw(embedding.weights[0]))
+
+        # 2D variable not excluded
+        var2 = backend.Variable([[1.0, 2.0], [3.0, 4.0]])
         optimizer = Muon()
-        self.assertAllClose(
-            False,
-            optimizer._should_use_adamw(vars),
-        )
+        self.assertFalse(optimizer._should_use_adamw(var2))
+
+        # Dense layer
         dense = Dense(2)
         dense.build([None, 2])
-        self.assertAllClose(
-            False,
-            optimizer._should_use_adamw(dense.weights[0]),
-        )
+        self.assertFalse(optimizer._should_use_adamw(dense.weights[0]))
+
+        # Dimension rules
+        v_1d = backend.Variable([1.0, 2.0], name="v1d")
+        v_5d = backend.Variable(np.zeros((2, 2, 2, 2, 2)), name="v5d")
+        self.assertTrue(optimizer._should_use_adamw(v_1d))
+        self.assertTrue(optimizer._should_use_adamw(v_5d))
 
     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])
-        optimizer._muon_update_step(grads, vars, 0.5)
-        self.assertAllClose(
-            vars, [[1.13, 1.51], [2.57, 4.06]], rtol=1e-2, atol=1e-2
-        )
+        var = backend.Variable([[1.0, 2.0], [3.0, 4.0]])
+        optimizer.build([var])
+        optimizer._muon_update_step(grads, var, 0.5)
+        self.assertAllClose(var, [[1.13, 1.51], [2.57, 4.06]], rtol=1e-2, atol=1e-2)
 
     def test_clip_norm(self):
         optimizer = Muon(clipnorm=1)
@@ -81,3 +74,13 @@ 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+ ResourceVariable (no .path)."""
+        optimizer = Muon()
+        var = backend.Variable([1.0, 2.0], name="test_var")
+        try:
+            result = optimizer._should_use_adamw(var)
+            self.assertIn(result, [True, False])
+        except AttributeError as e:
+            self.fail(f"Unexpected AttributeError: {e}")