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

harryjulian · Dec 21, 2023 · 1d053e0 · 1d053e0
1 parent 0aa65c7
commit 1d053e0
Showing 1 changed file with 9 additions and 10 deletions.
diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py
@@ -571,16 +571,15 @@ def test_rope(self):
             self.assertTrue(y.dtype, mx.float16)
 
     def test_alibi(self):
-        for kwargs in [{"num_heads": 8}]:
-            alibi = nn.ALibi(**kwargs)
-            shape = [1, 8, 20, 20]
-            x = mx.random.uniform(shape=shape)
-            y = alibi(x)
-            self.assertTrue(y.shape, shape)
-            self.assertTrue(y.dtype, mx.float32)
-
-            y = alibi(x.astype(mx.float16))
-            self.assertTrue(y.dtype, mx.float16)
+        alibi = nn.ALiBi()
+        shape = [1, 8, 20, 20]
+        x = mx.random.uniform(shape=shape)
+        y = alibi(x)
+        self.assertTrue(y.shape, shape)
+        self.assertTrue(y.dtype, mx.float32)
+
+        y = alibi(x.astype(mx.float16))
+        self.assertTrue(y.dtype, mx.float16)
 
 
 if __name__ == "__main__":