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

[Issue #1187] Add nan_to_num function initial attempt (#1247)

Browse Source 
* initial attempt, working with wrong types

* not compiling; mx.float16 and mx.bfloat16 tests added

* fix nan to num

* nit

---------

Co-authored-by: Awni Hannun <awni@apple.com>
This commit is contained in:
Anton Belov
2024-07-25 17:57:37 +01:00
committed by GitHub
parent baf9fa5f42
commit 5029894662
5 changed files with 93 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
python/tests/test_ops.py
View File

@@ -1653,6 +1653,23 @@ class TestOps(mlx_tests.MLXTestCase):
np.where,
)
def test_nan_to_num(self):
a = mx.array([6, float("inf"), 2, 0])
out_mx = mx.nan_to_num(a)
out_np = np.nan_to_num(a)
self.assertTrue(np.allclose(out_mx, out_np))
for t in [mx.float32, mx.float16]:
a = mx.array([float("inf"), 6.9, float("nan"), float("-inf")])
out_mx = mx.nan_to_num(a)
out_np = np.nan_to_num(a)
self.assertTrue(np.allclose(out_mx, out_np))
a = mx.array([float("inf"), 6.9, float("nan"), float("-inf")]).astype(t)
out_np = np.nan_to_num(a, nan=0.0, posinf=1000, neginf=-1000)
out_mx = mx.nan_to_num(a, nan=0.0, posinf=1000, neginf=-1000)
self.assertTrue(np.allclose(out_mx, out_np))
def test_as_strided(self):
x_npy = np.random.randn(128).astype(np.float32)
x_mlx = mx.array(x_npy)