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

Add mx.finfo and use it when making causal mask (#1726)

Browse Source 
* finfo

* fixes

* docs
This commit is contained in:
Awni Hannun
2024-12-19 14:52:41 -08:00
committed by GitHub
parent e03f0372b1
commit c3628eea49
9 changed files with 154 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
python/tests/test_nn.py
View File

@@ -1826,6 +1826,15 @@ class TestLayers(mlx_tests.MLXTestCase):
self.assertGreater(cosine(y, yq).min(), 0.99)
def test_causal_mask(self):
mask = nn.MultiHeadAttention.create_additive_causal_mask(4, mx.float16)
self.assertFalse(mx.any(mx.isnan(mask)))
self.assertTrue(mask[0, -1].item() < 0)
mask = nn.MultiHeadAttention.create_additive_causal_mask(4, mx.bfloat16)
self.assertFalse(mx.any(mx.isnan(mask)))
self.assertTrue(mask[0, -1].item() < 0)
if __name__ == "__main__":
unittest.main()