electron-rare/mlx-lm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix flaky test (#643)

Browse Source
This commit is contained in:
Awni Hannun
2025-12-02 13:33:11 -08:00
committed by GitHub
parent 34cbb8b51a
commit f940cf3a95
1 changed files with 4 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_losses.py
+4
View File
@@ -11,6 +11,7 @@ class TestLosses(unittest.TestCase):
def test_kl_div_loss(self):
self.assertTrue(can_run_metal())
mx.random.seed(0)
logits_q = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
logits_p = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
@@ -23,6 +24,7 @@ class TestLosses(unittest.TestCase):
def test_js_div_loss(self):
self.assertTrue(can_run_metal())
mx.random.seed(0)
logits_q = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
logits_p = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
@@ -35,6 +37,7 @@ class TestLosses(unittest.TestCase):
def test_kl_div_loss_vjp(self):
self.assertTrue(can_run_metal())
mx.random.seed(0)
logits_q = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
logits_p = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
@@ -48,6 +51,7 @@ class TestLosses(unittest.TestCase):
def test_js_div_loss_vjp(self):
self.assertTrue(can_run_metal())
mx.random.seed(0)
logits_q = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
logits_p = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)