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

mlx/ops.cpp

@@ -1,5 +1,4 @@
 // Copyright © 2023-2024 Apple Inc.
-
 #include <algorithm>
 #include <climits>
 #include <cmath>
@@ -1344,6 +1343,40 @@ array where(
       inputs);
 }
 
+array nan_to_num(
+    const array& a,
+    float nan /* = 0.0f */,
+    const std::optional<float>& posinf_ /* = std::nullopt */,
+    const std::optional<float>& neginf_ /* = std::nullopt */,
+    StreamOrDevice s /* = {} */) {
+  Dtype dtype = a.dtype();
+  if (!issubdtype(dtype, inexact)) {
+    return a;
+  }
+
+  auto type_to_max = [](const auto& dtype) -> float {
+    if (dtype == float32) {
+      return std::numeric_limits<float>::max();
+    } else if (dtype == bfloat16) {
+      return std::numeric_limits<bfloat16_t>::max();
+    } else if (dtype == float16) {
+      return std::numeric_limits<float16_t>::max();
+    } else {
+      std::ostringstream msg;
+      msg << "[nan_to_num] Does not yet support given type: " << dtype << ".";
+      throw std::invalid_argument(msg.str());
+    }
+  };
+
+  float posinf = posinf_ ? *posinf_ : type_to_max(dtype);
+  float neginf = neginf_ ? *neginf_ : -type_to_max(dtype);
+
+  auto out = where(isnan(a, s), array(nan, dtype), a, s);
+  out = where(isposinf(a, s), array(posinf, dtype), out, s);
+  out = where(isneginf(a, s), array(neginf, dtype), out, s);
+  return out;
+}
+
 array allclose(
     const array& a,
     const array& b,

mlx/ops.h

@@ -406,6 +406,14 @@ array where(
     const array& y,
     StreamOrDevice s = {});
 
+/** Replace NaN and infinities with finite numbers. */
+array nan_to_num(
+    const array& a,
+    float nan = 0.0f,
+    const std::optional<float>& posinf = std::nullopt,
+    const std::optional<float>& neginf = std::nullopt,
+    StreamOrDevice s = {});
+
 /** True if all elements in the array are true (or non-zero). **/
 array all(const array& a, bool keepdims, StreamOrDevice s = {});
 inline array all(const array& a, StreamOrDevice s = {}) {