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

Fix the hard-shrink test (#1185)

Browse Source
This commit is contained in:
Angelos Katharopoulos 2024-06-04 16:22:56 -07:00 committed by GitHub
parent 0b7d71fd2f
commit 0fe6895893
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
python/tests/test_nn.py
View File

@ -917,13 +917,13 @@ class TestLayers(mlx_tests.MLXTestCase):
def test_hard_shrink(self): def test_hard_shrink(self):
x = mx.array([1.0, -0.5, 0.0, 0.5, -1.5]) x = mx.array([1.0, -0.5, 0.0, 0.5, -1.5])
y = nn.hard_shrink(x) 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.assertTrue(mx.array_equal(y, expected_y))
self.assertEqual(y.shape, (5,)) self.assertEqual(y.shape, (5,))
self.assertEqual(y.dtype, mx.float32) self.assertEqual(y.dtype, mx.float32)
y = nn.hard_shrink(x, lambd=1.0) y = nn.hard_shrink(x, lambd=0.1)
expected_y = mx.array([1.0, 0.0, 0.0, 0.0, -1.5]) expected_y = mx.array([1.0, -0.5, 0.0, 0.5, -1.5])
self.assertTrue(mx.array_equal(y, expected_y)) self.assertTrue(mx.array_equal(y, expected_y))
self.assertEqual(y.shape, (5,)) self.assertEqual(y.shape, (5,))
self.assertEqual(y.dtype, mx.float32) self.assertEqual(y.dtype, mx.float32)