Add isnan (#423)

python/src/array.cpp

@@ -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,

python/src/ops.cpp

@@ -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,