From 508bd25e29c3cb7e6b759b949e303178f3153f53 Mon Sep 17 00:00:00 2001
From: Awni Hannun <awni@apple.com>
Date: Fri, 18 Jul 2025 06:43:11 -0700
Subject: [PATCH] match muon

---
 python/mlx/optimizers/optimizers.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/python/mlx/optimizers/optimizers.py b/python/mlx/optimizers/optimizers.py
index 71d2ab9ba..07b68cc5b 100644
--- a/python/mlx/optimizers/optimizers.py
+++ b/python/mlx/optimizers/optimizers.py
@@ -903,8 +903,7 @@ class Muon(Optimizer):
         if transpose_needed:
             X = X.T
 
-        norm = mx.sqrt(mx.sum(mx.square(X), keepdims=True) + 1e-7)
-        X = X / norm
+        X = X / (mx.linalg.norm(X, keepdims=True) + 1e-7)
 
         for _ in range(steps):
             A = X @ X.T