zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-06-24 09:21:18 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Use mx.allclose

Browse Source
This commit is contained in:
Neil Mehta 2025-03-11 10:24:26 -04:00
parent e6031a78e4
commit 21da261028
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
llms/tests/test_sample_utils.py
View File

@ -15,7 +15,7 @@ class TestSampleUtils(unittest.TestCase):
new_logits = apply_top_p(logits, 0.95) new_logits = apply_top_p(logits, 0.95)
actual_probs = mx.softmax(new_logits.squeeze()) actual_probs = mx.softmax(new_logits.squeeze())
self.assertEqual(probs.squeeze().tolist(), actual_probs.tolist()) self.assertTrue(mx.allclose(probs.squeeze(), actual_probs))
probs = mx.array([0.0, 0.5, 0.4, 0.1])[None] probs = mx.array([0.0, 0.5, 0.4, 0.1])[None]
logits = mx.log(probs) logits = mx.log(probs)
@ -56,7 +56,7 @@ class TestSampleUtils(unittest.TestCase):
logits = mx.log(probs) logits = mx.log(probs)
new_logits = apply_min_p(logits, 0.05) new_logits = apply_min_p(logits, 0.05)
actual_probs = mx.softmax(new_logits.squeeze()) actual_probs = mx.softmax(new_logits.squeeze())
self.assertEqual(actual_probs.tolist(), mx.squeeze(probs).tolist()) self.assertTrue(mx.allclose(actual_probs, mx.squeeze(probs)))
# Batch mode works # Batch mode works
probs = mx.array([[0.9, 0.0, 0.0, 0.1], [0.0, 0.8, 0.0, 0.1]]) probs = mx.array([[0.9, 0.0, 0.0, 0.1], [0.0, 0.8, 0.0, 0.1]])