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

implemented isposinf and isneginf in one PR (#470)

Browse Source 
* ran precommit

* updated docs
This commit is contained in:
Yashraj Singh
2024-01-16 20:18:07 +05:30
committed by GitHub
parent a2ffea683a
commit e72458a3fa
6 changed files with 128 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
python/tests/test_ops.py
View File

@@ -401,6 +401,36 @@ class TestOps(mlx_tests.MLXTestCase):
with self.assertRaises(ValueError):
mx.ceil(mx.array([22 + 3j, 19 + 98j]))
def test_isposinf(self):
x = mx.array([0.0, float("-inf")])
self.assertEqual(mx.isposinf(x).tolist(), [False, False])
x = mx.array([0.0, float("-inf")]).astype(mx.float16)
self.assertEqual(mx.isposinf(x).tolist(), [False, False])
x = mx.array([0.0, float("-inf")]).astype(mx.bfloat16)
self.assertEqual(mx.isposinf(x).tolist(), [False, False])
x = mx.array([0.0, float("-inf")]).astype(mx.complex64)
self.assertEqual(mx.isposinf(x).tolist(), [False, False])
self.assertEqual(mx.isposinf(0 * mx.array(float("inf"))).tolist(), False)
def test_isneginf(self):
x = mx.array([0.0, float("-inf")])
self.assertEqual(mx.isneginf(x).tolist(), [False, True])
x = mx.array([0.0, float("-inf")]).astype(mx.float16)
self.assertEqual(mx.isneginf(x).tolist(), [False, True])
x = mx.array([0.0, float("-inf")]).astype(mx.bfloat16)
self.assertEqual(mx.isneginf(x).tolist(), [False, True])
x = mx.array([0.0, float("-inf")]).astype(mx.complex64)
self.assertEqual(mx.isneginf(x).tolist(), [False, True])
self.assertEqual(mx.isneginf(0 * mx.array(float("inf"))).tolist(), False)
def test_round(self):
# float
x = mx.array(