Repository: keras-team/keras
Pull request #21079: fix muon document
Author: login: pass-lin
Labels: name: kokoro:force-run
name: ready to pull
description: Ready to be merged into the codebase
name: size:XS
Created: 2025-03-21T12:40:40Z | Merged: 2025-03-23T21:10:22Z

I'm the author of [#21037](https://github.com/keras-team/keras/commit/717a64a3a467d4ee19dd9314e8e2402593d826bd), and I found that I made a mistake in this document. I'll just revise it again.


Diff:
diff --git a/keras/src/optimizers/muon.py b/keras/src/optimizers/muon.py
index c74cfa8cf1c5..853424ca2ae6 100644
--- a/keras/src/optimizers/muon.py
+++ b/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):
diff --git a/keras/src/optimizers/muon_test.py b/keras/src/optimizers/muon_test.py
new file mode 100644
index 000000000000..f22423c34aae
--- /dev/null
+++ b/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])


Reviews:
user: login: fchollet
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-21T16:58:12Z
user: login: fchollet
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-21T22:33:36Z
user: login: fchollet
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-21T22:33:58Z
user: login: fchollet
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-22T23:10:05Z
user: login: fchollet
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-22T23:14:41Z
user: login: fchollet
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-23T21:10:17Z

Review comments:
diff_hunk: @@ -53,7 +53,7 @@ class Muon(optimizer.Optimizer):
             (in the formula just before Section 2.1),
             not the epsilon in Algorithm 1 of the paper.
             It be used at Adamw.Defaults to `1e-7`.
-        exclude_layers: List of strings, keywords of layer names to exclude.
+        exclude_layers: str, keywords of layer names to exclude.
path: keras/src/optimizers/muon.py
user: login: fchollet
body: I believe we should change the *behavior* of the argument to make it a list of strs. The most intuitive usage would be to pass the list of layer names you want to exclude.
created_at: 2025-03-21T16:58:12Z
author_association: COLLABORATOR
diff_hunk: @@ -53,7 +53,7 @@ class Muon(optimizer.Optimizer):
             (in the formula just before Section 2.1),
             not the epsilon in Algorithm 1 of the paper.
             It be used at Adamw.Defaults to `1e-7`.
-        exclude_layers: List of strings, keywords of layer names to exclude.
+        exclude_layers: List of strings,, keywords of layer names to exclude.
path: keras/src/optimizers/muon.py
user: login: fchollet
body: Please remove double ","
created_at: 2025-03-21T22:33:36Z
author_association: COLLABORATOR
diff_hunk: @@ -88,7 +88,7 @@ def __init__(
         loss_scale_factor=None,
         gradient_accumulation_steps=None,
         name="muon",
-        exclude_layers=None,
+        exclude_layers=[],
path: keras/src/optimizers/muon.py
user: login: fchollet
body: Initialize with None. Then do `exclude_layers = exclude_layers or []` in the constructor to turn None into empty list
created_at: 2025-03-21T22:33:58Z
author_association: COLLABORATOR
diff_hunk: @@ -124,10 +124,9 @@ 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()
+        if exclude_layers is None:
+            exclude_layers = []
+        self.exclude_layers = exclude_layers
path: keras/src/optimizers/muon.py
user: login: fchollet
body: Not very important, but these 3 lines are equivalent to one: `self.exclude_layers = exclude_layers or []`
created_at: 2025-03-22T23:10:05Z
author_association: COLLABORATOR
diff_hunk: @@ -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 keyword.lower() in variable.path.lower():
path: keras/src/optimizers/muon.py
user: login: fchollet
body: Some inefficiency here -- use `if re.search(keyword, variable.path):`
created_at: 2025-03-22T23:14:40Z
author_association: COLLABORATOR

PR comments:
user: login: codecov-commenter
created_at: 2025-03-21T12:45:54Z
body: ## [Codecov](https://app.codecov.io/gh/keras-team/keras/pull/21079?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team) Report
Attention: Patch coverage is `80.00000%` with `1 line` in your changes missing coverage. Please review.
> Project coverage is 82.65%. Comparing base [(`9fea22a`)](https://app.codecov.io/gh/keras-team/keras/commit/9fea22aaca4ad23eafa46c4cf82700c61d5a88e5?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team) to head [(`9e67b7f`)](https://app.codecov.io/gh/keras-team/keras/commit/9e67b7fa287754e8ecc6bdfa2cd0248f0548ea7d?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team).
> Report is 43 commits behind head on master.

| [Files with missing lines](https://app.codecov.io/gh/keras-team/keras/pull/21079?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team) | Patch % | Lines |
|---|---|---|
| [keras/src/optimizers/muon.py](https://app.codecov.io/gh/keras-team/keras/pull/21079?src=pr&el=tree&filepath=keras%2Fsrc%2Foptimizers%2Fmuon.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team#diff-a2VyYXMvc3JjL29wdGltaXplcnMvbXVvbi5weQ==) | 80.00% | [0 Missing and 1 partial :warning: ](https://app.codecov.io/gh/keras-team/keras/pull/21079?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master   #21079      +/-   ##
==========================================
+ Coverage   82.54%   82.65%   +0.11%     
==========================================
  Files         564      564              
  Lines       54092    54097       +5     
  Branches     8403     8406       +3     
==========================================
+ Hits        44648    44716      +68     
+ Misses       7371     7303      -68     
- Partials     2073     2078       +5     
```

| [Flag](https://app.codecov.io/gh/keras-team/keras/pull/21079/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team) | Coverage Δ | |
|---|---|---|
| [keras](https://app.codecov.io/gh/keras-team/keras/pull/21079/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team) | `82.47% <80.00%> (+0.11%)` | :arrow_up: |
| [keras-jax](https://app.codecov.io/gh/keras-team/keras/pull/21079/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team) | `64.04% <80.00%> (+0.11%)` | :arrow_up: |
| [keras-numpy](https://app.codecov.io/gh/keras-team/keras/pull/21079/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team) | `59.12% <80.00%> (+0.11%)` | :arrow_up: |
| [keras-openvino](https://app.codecov.io/gh/keras-team/keras/pull/21079/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team) | `32.63% <20.00%> (-0.01%)` | :arrow_down: |
| [keras-tensorflow](https://app.codecov.io/gh/keras-team/keras/pull/21079/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team) | `64.36% <80.00%> (+0.11%)` | :arrow_up: |
| [keras-torch](https://app.codecov.io/gh/keras-team/keras/pull/21079/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team) | `64.09% <80.00%> (+0.11%)` | :arrow_up: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team#carryforward-flags-in-the-pull-request-comment) to find out more.

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/keras-team/keras/pull/21079?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE
user: login: pass-lin
created_at: 2025-03-21T20:24:24Z
body: @fchollet 
I tried to add some test cases to muon. The main purpose here is to facilitate verification of whether it can be used.
author_association: CONTRIBUTOR