modify code .

Author: pass-lin

keras/src/backend/tensorflow/optimizer.py

@@ -113,8 +113,9 @@ def weight_decay_fn(variable):
     def _backend_update_step(self, grads, trainable_variables, learning_rate):
         def _prepare_var(v):
             new_v = v.value if isinstance(v, backend.Variable) else v
-            new_v._muon_use_adam_flag = v._muon_use_adam_flag
-            new_v._muon_path_id = v._muon_path_id
+            if hasattr(v, "_muon_use_adam_flag"):
+                new_v._muon_use_adam_flag = v._muon_use_adam_flag
+                new_v._muon_path_id = v._muon_path_id
             return new_v
 
         trainable_variables = [_prepare_var(v) for v in trainable_variables]

keras/src/optimizers/muon_test.py

@@ -92,6 +92,7 @@ def test_muon_weight_decay(self):
         weight_decay = 0.01
         expected_variable = variable - variable * weight_decay
         optimizer = Muon(learning_rate=1.0, weight_decay=weight_decay)
+        optimizer.build([variable])
         optimizer._apply_weight_decay([variable])
         self.assertAllClose(variable, expected_variable, rtol=1e-4, atol=1e-4)
 
@@ -100,6 +101,7 @@ def test_adamw_weight_decay(self):
         weight_decay = 0.01
         expected_variable = variable - variable * weight_decay
         optimizer = Muon(learning_rate=1.0, adam_weight_decay=weight_decay)
+        optimizer.build([variable])
         optimizer._apply_weight_decay([variable])
 
         self.assertAllClose(variable, expected_variable, rtol=1e-4, atol=1e-4)