zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-19 08:38:09 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix the alibi test that was left unchanged (#252)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2023-12-21 14:59:25 -08:00
committed by GitHub
parent 0aa65c7a6b
commit 1d053e0d1d
1 changed files with 8 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
python/tests/test_nn.py
View File

@@ -571,16 +571,15 @@ class TestNN(mlx_tests.MLXTestCase):
self.assertTrue(y.dtype, mx.float16) self.assertTrue(y.dtype, mx.float16)
def test_alibi(self): def test_alibi(self):
for kwargs in [{"num_heads": 8}]: alibi = nn.ALiBi()
alibi = nn.ALibi(**kwargs) shape = [1, 8, 20, 20]
shape = [1, 8, 20, 20] x = mx.random.uniform(shape=shape)
x = mx.random.uniform(shape=shape) y = alibi(x)
y = alibi(x) self.assertTrue(y.shape, shape)
self.assertTrue(y.shape, shape) self.assertTrue(y.dtype, mx.float32)
self.assertTrue(y.dtype, mx.float32)
y = alibi(x.astype(mx.float16)) y = alibi(x.astype(mx.float16))
self.assertTrue(y.dtype, mx.float16) self.assertTrue(y.dtype, mx.float16)
if __name__ == "__main__": if __name__ == "__main__":