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

G.astype(mx.bfloat16) to G.astype(G.dtype)

Browse Source
This commit is contained in:
Goekdeniz-Guelmez
2025-07-17 19:49:26 +02:00
parent 7f39e9c299
commit 060404d862
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -896,7 +896,7 @@ class Muon(Optimizer):
def _zeropower_via_newtonschulz5(self, G, steps: int): def _zeropower_via_newtonschulz5(self, G, steps: int):
assert G.ndim >= 2 assert G.ndim >= 2
a, b, c = (3.4445, -4.7750, 2.0315) a, b, c = (3.4445, -4.7750, 2.0315)
X = G.astype(mx.bfloat16) X = G.astype(G.dtype)
transpose_needed = G.shape[-2] > G.shape[-1] transpose_needed = G.shape[-2] > G.shape[-1]
if transpose_needed: if transpose_needed: