From 698daee214e05c15e3aa21379e4f08f06fe588b2 Mon Sep 17 00:00:00 2001
From: Goekdeniz-Guelmez <Gulmezdeniz1999@gmail.com>
Date: Thu, 17 Jul 2025 19:57:18 +0200
Subject: [PATCH] replace with  mx.addmm

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

diff --git a/python/mlx/optimizers/optimizers.py b/python/mlx/optimizers/optimizers.py
index 77faffcc9..88820dbd3 100644
--- a/python/mlx/optimizers/optimizers.py
+++ b/python/mlx/optimizers/optimizers.py
@@ -909,8 +909,8 @@ class Muon(Optimizer):
         # Perform the NS iterations
         for _ in range(steps):
             A = X @ X.T
-            B = b * A + c * (A @ A)
-            X = a * X + B @ X
+            B = mx.addmm(b * A, A, A, beta=1.0, alpha=c)
+            X = mx.addmm(a * X, B, X, beta=1.0, alpha=1.0)
 
         if transpose_needed:
             X = X.T