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

* 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>

python/src/ops.cpp

@@ -3595,6 +3595,39 @@ void init_ops(nb::module_& m) {
             array: The output containing elements selected from
             ``x`` and ``y``.
       )pbdoc");
+  m.def(
+      "nan_to_num",
+      [](const ScalarOrArray& a,
+         float nan,
+         std::optional<float>& posinf,
+         std::optional<float>& neginf,
+         StreamOrDevice s) {
+        return nan_to_num(to_array(a), nan, posinf, neginf, s);
+      },
+      nb::arg(),
+      "nan"_a = 0.0f,
+      "posinf"_a = nb::none(),
+      "neginf"_a = nb::none(),
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def nan_to_num(a: Union[scalar, array], nan: float = 0, posinf: Optional[float] = None, neginf: Optional[float] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        Replace NaN and Inf values with finite numbers.
+
+        Args:
+            a (array): Input array
+            nan (float, optional): Value to replace NaN with. Default: ``0``.
+            posinf (float, optional): Value to replace positive infinities
+              with. If ``None``, defaults to largest finite value for the
+              given data type. Default: ``None``.
+            neginf (float, optional): Value to replace negative infinities
+              with. If ``None``, defaults to the negative of the largest
+              finite value for the given data type. Default: ``None``.
+
+        Returns:
+            array: Output array with NaN and Inf replaced.
+    )pbdoc");
   m.def(
       "round",
       [](const ScalarOrArray& a, int decimals, StreamOrDevice s) {

python/tests/test_ops.py

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