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

Fix flaky test (#1020)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2026-03-18 12:51:22 -07:00
committed by GitHub
parent 564281f793
commit f8019f7769
2 changed files with 12 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/pull_request.yml
+1 -1
View File
@@ -40,5 +40,5 @@ jobs:
run: |
curl -o test_data.zip -L https://github.com/ml-explore/mlx-lm/releases/download/test_data/test_data.zip
unzip test_data.zip
HF_HOME="." python -m xmlrunner discover -v tests -o test-results/
METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 HF_HOME="." python -m xmlrunner discover -v tests -o test-results/
mlx.launch -n 2 tests/model_parallel_tests.py
tests/test_losses.py
+11 -11
View File
@@ -13,21 +13,21 @@ class TestLosses(unittest.TestCase):
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)
logits_q = mx.random.normal((2, 4, 4000))
logits_p = mx.random.normal((2, 4, 4000))
with mx.stream(mx.cpu):
expected = kl_div_loss(logits_q, logits_p)
kl = kl_div_loss(logits_q, logits_p)
self.assertTrue(mx.allclose(kl, expected, rtol=1e-4))
self.assertTrue(mx.allclose(kl, expected))
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)
logits_q = mx.random.normal((2, 4, 4000))
logits_p = mx.random.normal((2, 4, 4000))
with mx.stream(mx.cpu):
expected = js_div_loss(logits_q, logits_p)
@@ -39,9 +39,9 @@ class TestLosses(unittest.TestCase):
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)
cotan = mx.random.uniform(shape=(4, 8), dtype=mx.float32)
logits_q = mx.random.normal((2, 4, 4000))
logits_p = mx.random.normal((2, 4, 4000))
cotan = mx.random.normal((2, 4))
with mx.stream(mx.cpu):
expected = mx.vjp(kl_div_loss, [logits_q, logits_p], [cotan])[1][0]
@@ -53,9 +53,9 @@ class TestLosses(unittest.TestCase):
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)
cotan = mx.random.uniform(shape=(4, 8), dtype=mx.float32)
logits_q = mx.random.normal((2, 4, 4000))
logits_p = mx.random.normal((2, 4, 4000))
cotan = mx.random.normal((2, 4))
with mx.stream(mx.cpu):
expected = mx.vjp(js_div_loss, [logits_q, logits_p], [cotan])[1][0]