zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-11-01 16:58:08 +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

12
python/tests/test_array.py
View File

@@ -97,6 +97,18 @@ class TestDtypes(mlx_tests.MLXTestCase):
self.assertListEqual(list(z.shape), list(x.shape))
self.assertListEqual(list(z.shape), list(y.shape))
def test_finfo(self):
with self.assertRaises(ValueError):
mx.finfo(mx.int32)
self.assertEqual(mx.finfo(mx.float32).min, np.finfo(np.float32).min)
self.assertEqual(mx.finfo(mx.float32).max, np.finfo(np.float32).max)
self.assertEqual(mx.finfo(mx.float32).dtype, mx.float32)
self.assertEqual(mx.finfo(mx.float16).min, np.finfo(np.float16).min)
self.assertEqual(mx.finfo(mx.float16).max, np.finfo(np.float16).max)
self.assertEqual(mx.finfo(mx.float16).dtype, mx.float16)
class TestEquality(mlx_tests.MLXTestCase):
def test_array_eq_array(self):