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

Author: utsab345

keras/src/optimizers/muon.py

@@ -128,19 +128,20 @@ 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 
+        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:
+        if not 1 < len(variable.shape) < 5:
             return True
 
         # Exclude embedding layers if specified
-        var_identifier = getattr(variable, "name", "") or getattr(variable, "path", "")
+        var_identifier = getattr(variable, "name", "") or getattr(
+            variable, "path", ""
+        )
         if self.exclude_embeddings and "embedding" in var_identifier.lower():
             return True
 
@@ -156,7 +157,6 @@ def _should_use_adamw(self, variable):
         # Otherwise, use AdamW
         return False
 
-
     def build(self, var_list):
         """Initialize optimizer variables.
 

keras/src/optimizers/muon_test.py

@@ -3,7 +3,8 @@
 from keras.src import backend
 from keras.src import ops
 from keras.src import testing
-from keras.src.layers import Dense, Embedding
+from keras.src.layers import Dense
+from keras.src.layers import Embedding
 from keras.src.optimizers.muon import Muon
 
 
@@ -61,7 +62,9 @@ def test_muon_single_step(self):
         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)
+        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)
@@ -76,7 +79,8 @@ def test_clip_value(self):
         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)."""
+        """Ensure compatibility with TF 2.16+
+        ResourceVariable (no .path)."""
         optimizer = Muon()
         var = backend.Variable([1.0, 2.0], name="test_var")
         try: