zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

match muon

Browse Source
This commit is contained in:
Awni Hannun
2025-07-18 06:43:11 -07:00
parent 0a8bb904d7
commit 508bd25e29
1 changed files with 1 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/mlx/optimizers/optimizers.py
View File

@@ -903,8 +903,7 @@ class Muon(Optimizer):
if transpose_needed: if transpose_needed:
X = X.T X = X.T
norm = mx.sqrt(mx.sum(mx.square(X), keepdims=True) + 1e-7) X = X / (mx.linalg.norm(X, keepdims=True) + 1e-7)
X = X / norm
for _ in range(steps): for _ in range(steps):
A = X @ X.T A = X @ X.T