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

Add isnan (#423)

Browse Source
This commit is contained in:
Ayush Shridhar
2024-01-12 14:16:48 -05:00
committed by GitHub
parent 29081204d1
commit 1416e7b664
7 changed files with 72 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/src/array.cpp
View File

@@ -1132,6 +1132,12 @@ void init_array(py::module_& m) {
py::kw_only(),
"stream"_a = none,
"See :func:`any`.")
.def(
"isnan",
&mlx::core::isnan,
py::kw_only(),
"stream"_a = none,
"See :func:`isnan`.")
.def(
"moveaxis",
&moveaxis,

18
python/src/ops.cpp
View File

@@ -1820,6 +1820,24 @@ void init_ops(py::module_& m) {
Returns:
array: The ceil of ``a``.
)pbdoc");
m.def(
"isnan",
&mlx::core::isnan,
"a"_a,
py::pos_only(),
py::kw_only(),
"stream"_a = none,
R"pbdoc(
isnan(a: array, stream: Union[None, Stream, Device] = None) -> array
Return a boolean array indicating which elements are NaN.
Args:
a (array): Input array.
Returns:
array: The array with boolean values indicating which elements are NaN.
)pbdoc");
m.def(
"moveaxis",
&moveaxis,

15
python/tests/test_ops.py
View File

@@ -321,6 +321,21 @@ class TestOps(mlx_tests.MLXTestCase):
self.assertFalse(mx.array_equal(x, y))
self.assertTrue(mx.array_equal(x, y, equal_nan=True))
def test_isnan(self):
x = mx.array([0.0, float("nan")])
self.assertEqual(mx.isnan(x).tolist(), [False, True])
x = mx.array([0.0, float("nan")]).astype(mx.float16)
self.assertEqual(mx.isnan(x).tolist(), [False, True])
x = mx.array([0.0, float("nan")]).astype(mx.bfloat16)
self.assertEqual(mx.isnan(x).tolist(), [False, True])
x = mx.array([0.0, float("nan")]).astype(mx.complex64)
self.assertEqual(mx.isnan(x).tolist(), [False, True])
self.assertEqual(mx.isnan(0 * mx.array(float("inf"))).tolist(), True)
def test_tri(self):
for shape in [[4], [4, 4], [2, 10]]:
for diag in [-1, 0, 1, -2]: