modify code .

Author: pass-lin

keras/src/backend/tensorflow/optimizer.py

@@ -114,7 +114,7 @@ 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
+            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.py

@@ -171,16 +171,16 @@ def build(self, var_list):
         self.muon_velocities = {}
 
         for var in var_list:
-            var.muon_path_id = self._var_key(var)
+            var._muon_path_id = self._var_key(var)
             if not self._overwrite_variable_with_gradient(var):
-                self.adam_momentums[var.muon_path_id] = (
+                self.adam_momentums[var._muon_path_id] = (
                     self.add_variable_from_reference(
                         reference_variable=var, name="momentum"
                     )
                 )
                 var._muon_use_adam_flag = self._should_use_adamw(var)
                 if var._muon_use_adam_flag:
-                    self.adam_velocities[var.muon_path_id] = (
+                    self.adam_velocities[var._muon_path_id] = (
                         self.add_variable_from_reference(
                             reference_variable=var, name="velocity"
                         )
@@ -196,7 +196,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.muon_path_id]
+        m = self.adam_momentums[variable._muon_path_id]
         self.assign_add(m, ops.add(gradient, m * (self.momentum - 1)))
         if self.nesterov:
             g = ops.add(gradient, self.momentum * m)
@@ -221,8 +221,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.muon_path_id]
-        v = self.adam_velocities[variable.muon_path_id]
+        m = self.adam_momentums[variable._muon_path_id]
+        v = self.adam_velocities[variable._muon_path_id]
 
         alpha = lr * ops.sqrt(1 - adam_beta_2_power) / (1 - adam_beta_1_power)