From 21da261028f4f0d65e51f5f6452a314ec48ea041 Mon Sep 17 00:00:00 2001 From: Neil Mehta Date: Tue, 11 Mar 2025 10:24:26 -0400 Subject: [PATCH] Use mx.allclose --- llms/tests/test_sample_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/llms/tests/test_sample_utils.py b/llms/tests/test_sample_utils.py index a8664fd9..7760c569 100644 --- a/llms/tests/test_sample_utils.py +++ b/llms/tests/test_sample_utils.py @@ -15,7 +15,7 @@ class TestSampleUtils(unittest.TestCase): new_logits = apply_top_p(logits, 0.95) 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] logits = mx.log(probs) @@ -56,7 +56,7 @@ class TestSampleUtils(unittest.TestCase): logits = mx.log(probs) new_logits = apply_min_p(logits, 0.05) 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 probs = mx.array([[0.9, 0.0, 0.0, 0.1], [0.0, 0.8, 0.0, 0.1]])