From 0fe68958935286dbc3950945c188d1885234fe23 Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Tue, 4 Jun 2024 16:22:56 -0700 Subject: [PATCH] Fix the hard-shrink test (#1185) --- python/tests/test_nn.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index 27537cac4..70b2eecc4 100644 --- a/python/tests/test_nn.py +++ b/python/tests/test_nn.py @@ -917,13 +917,13 @@ class TestLayers(mlx_tests.MLXTestCase): def test_hard_shrink(self): x = mx.array([1.0, -0.5, 0.0, 0.5, -1.5]) y = nn.hard_shrink(x) - expected_y = mx.array([1.0, -0.5, 0.0, 0.5, -1.5]) + expected_y = mx.array([1.0, 0.0, 0.0, 0.0, -1.5]) self.assertTrue(mx.array_equal(y, expected_y)) self.assertEqual(y.shape, (5,)) self.assertEqual(y.dtype, mx.float32) - y = nn.hard_shrink(x, lambd=1.0) - expected_y = mx.array([1.0, 0.0, 0.0, 0.0, -1.5]) + y = nn.hard_shrink(x, lambd=0.1) + expected_y = mx.array([1.0, -0.5, 0.0, 0.5, -1.5]) self.assertTrue(mx.array_equal(y, expected_y)) self.assertEqual(y.shape, (5,)) self.assertEqual(y.dtype, mx.float32)