fix muon document (#21079)

* fix muon argument

* fix muon argument

* change behavior

* add some test

* add some test

* fix

* fix

Author: pass-lin

keras/src/optimizers/muon.py

@@ -1,3 +1,5 @@
+import re
+
 from keras.src import ops
 from keras.src.api_export import keras_export
 from keras.src.optimizers import optimizer
@@ -124,10 +126,7 @@ def __init__(
         self.ns_steps = ns_steps
         self.nesterov = nesterov
         self.exclude_embeddings = exclude_embeddings
-        # exclude_layers is a keyword at variable path
-        # so it must be a string
-        assert isinstance(exclude_layers, str) or exclude_layers is None
-        self.exclude_layers = exclude_layers.lower()
+        self.exclude_layers = exclude_layers or []
 
     def _should_use_adamw(self, variable):
         # To use it with 4D convolutional filters,
@@ -137,8 +136,9 @@ def _should_use_adamw(self, variable):
             return True
         if self.exclude_embeddings and "embedding" in variable.path.lower():
             return True
-        if self.exclude_layers in variable.path.lower():
-            return True
+        for keyword in self.exclude_layers:
+            if re.search(keyword, variable.path):
+                return True
         return False
 
     def build(self, var_list):

keras/src/optimizers/muon_test.py

@@ -0,0 +1,83 @@
+import numpy as np
+
+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.optimizers.muon import Muon
+
+
+class MuonTest(testing.TestCase):
+    def test_config(self):
+        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]])
+        output = optimizer.zeropower_via_newtonschulz5(tensor_input, 5)
+        self.assertAllClose(output, except_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)
+
+    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),
+        )
+        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]])
+        optimizer = Muon()
+        self.assertAllClose(
+            False,
+            optimizer._should_use_adamw(vars),
+        )
+        dense = Dense(2)
+        dense.build([None, 2])
+        self.assertAllClose(
+            False,
+            optimizer._should_use_adamw(dense.weights[0]),
+        )
+
+    def test_muon_single_step(self):
+        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
+        )
+
+    def test_clip_norm(self):
+        optimizer = Muon(clipnorm=1)
+        grad = [np.array([100.0, 100.0])]
+        clipped_grad = optimizer._clip_gradients(grad)
+        self.assertAllClose(clipped_grad[0], [2**0.5 / 2, 2**0.5 / 2])
+
+    def test_clip_value(self):
+        optimizer = Muon(clipvalue=1)
+        grad = [np.array([100.0, 100.0])]
+        clipped_grad = optimizer._clip_gradients(grad)
+        self.assertAllClose(clipped_grad[0], [1.0, 1.0])