zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-18 07:18:29 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix type promotion in Adam with bias correction (#2350)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2025-07-10 11:14:42 -07:00
committed by GitHub
parent afb9817599
commit 0eb035b4b1
2 changed files with 11 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
python/tests/test_optimizers.py
View File

@@ -196,6 +196,13 @@ class TestOptimizers(mlx_tests.MLXTestCase):
)
)
# Test for correct gradient type propagation
params = tree_map(lambda x: x.astype(mx.float16), params)
grads = tree_map(lambda x: x.astype(mx.float16), grads)
optim = opt.Adam(1e-2, bias_correction=True)
new_params = optim.apply_gradients(grads, params)
self.assertTrue(tree_equal(lambda p: p.dtype == mx.float16, new_params))
@unittest.skipIf(not has_torch, "requires Torch")
def test_adamw_matches_pytorch(self):
mx.random.seed(0)