iinfo and scalar overflow detection (#2009)

python/src/array.cpp

@@ -206,6 +206,30 @@ void init_array(nb::module_& m) {
         return os.str();
       });
 
+  nb::class_<mx::iinfo>(
+      m,
+      "iinfo",
+      R"pbdoc(
+      Get information on integer types.
+      )pbdoc")
+      .def(nb::init<mx::Dtype>())
+      .def_ro(
+          "min",
+          &mx::iinfo::min,
+          R"pbdoc(The smallest representable number.)pbdoc")
+      .def_ro(
+          "max",
+          &mx::iinfo::max,
+          R"pbdoc(The largest representable number.)pbdoc")
+      .def_ro("dtype", &mx::iinfo::dtype, R"pbdoc(The :obj:`Dtype`.)pbdoc")
+      .def("__repr__", [](const mx::iinfo& i) {
+        std::ostringstream os;
+        os << "iinfo("
+           << "min=" << i.min << ", max=" << i.max << ", dtype=" << i.dtype
+           << ")";
+        return os.str();
+      });
+
   nb::class_<ArrayAt>(
       m,
       "ArrayAt",

python/src/utils.cpp

@@ -2,6 +2,7 @@
 
 #include "python/src/utils.h"
 #include "mlx/ops.h"
+#include "mlx/utils.h"
 #include "python/src/convert.h"
 
 mx::array to_array(
@@ -16,6 +17,16 @@ mx::array to_array(
         ? mx::int64
         : mx::int32;
     auto out_t = dtype.value_or(default_type);
+    if (mx::issubdtype(out_t, mx::integer) && out_t.size() < 8) {
+      auto info = mx::iinfo(out_t);
+      if (val < info.min || val > static_cast<int64_t>(info.max)) {
+        std::ostringstream msg;
+        msg << "Converting " << val << " to " << out_t
+            << " would result in overflow.";
+        throw std::invalid_argument(msg.str());
+      }
+    }
+
     // bool_ is an exception and is always promoted
     return mx::array(val, (out_t == mx::bool_) ? mx::int32 : out_t);
   } else if (auto pv = std::get_if<nb::float_>(&v); pv) {