python: Widen int scalars that do not fit in int64 to uint64

python/src/convert.cpp

@@ -557,7 +557,12 @@ PyScalarT validate_shape(
         t = pyint;
         // Match the scalar path, which widens to int64 rather than failing
         // when a python int does not fit in int32.
-        auto val = nb::cast<int64_t>(l);
+        auto [val, is_uint64] = to_int64_or_uint64(l);
+        if (is_uint64) {
+          throw std::invalid_argument(
+              "Python ints larger than int64 are not supported in list "
+              "initialization.");
+        }
         if (val > std::numeric_limits<int>::max() ||
             val < std::numeric_limits<int>::min()) {
           has_wide_int = true;
@@ -730,7 +735,10 @@ mx::array create_array(
   if (nb::isinstance<nb::bool_>(v)) {
     return mx::array(nb::cast<bool>(v), t.value_or(mx::bool_));
   } else if (nb::isinstance<nb::int_>(v)) {
-    auto val = nb::cast<int64_t>(v);
+    auto [val, is_uint64] = to_int64_or_uint64(v);
+    if (is_uint64) {
+      return mx::array(static_cast<uint64_t>(val), t.value_or(mx::uint64));
+    }
     auto default_type = (val > std::numeric_limits<int>::max() ||
                          val < std::numeric_limits<int>::min())
         ? mx::int64

python/src/utils.cpp

@@ -11,7 +11,18 @@ mx::array to_array(
   if (auto pv = std::get_if<nb::bool_>(&v); pv) {
     return mx::array(nb::cast<bool>(*pv), dtype.value_or(mx::bool_));
   } else if (auto pv = std::get_if<nb::int_>(&v); pv) {
-    auto val = nb::cast<int64_t>(*pv);
+    auto [val, is_uint64] = to_int64_or_uint64(*pv);
+    if (is_uint64) {
+      auto out_t = dtype.value_or(mx::uint64);
+      if (mx::issubdtype(out_t, mx::integer) && out_t != mx::uint64) {
+        std::ostringstream msg;
+        msg << "Converting " << static_cast<uint64_t>(val) << " to " << out_t
+            << " would result in overflow.";
+        throw std::invalid_argument(msg.str());
+      }
+      return mx::array(
+          static_cast<uint64_t>(val), (out_t == mx::bool_) ? mx::int32 : out_t);
+    }
     auto default_type = (val > std::numeric_limits<int>::max() ||
                          val < std::numeric_limits<int>::min())
         ? mx::int64

python/src/utils.h

@@ -67,6 +67,24 @@ inline void throw_invalid_operation(
   throw std::invalid_argument(msg.str());
 }
 
+// Parse a python int as int64, or as uint64 (flagged true) when it only
+// fits there. Raise a clean error when neither type can hold it.
+inline std::pair<int64_t, bool> to_int64_or_uint64(nb::handle h) {
+  int64_t val;
+  if (nb::try_cast<int64_t>(h, val)) {
+    return {val, false};
+  }
+  uint64_t uval = PyLong_AsUnsignedLongLong(h.ptr());
+  if (!PyErr_Occurred()) {
+    return {static_cast<int64_t>(uval), true};
+  }
+  PyErr_Clear();
+  std::ostringstream msg;
+  msg << "Python int " << nb::str(h).c_str()
+      << " does not fit in int64 or uint64.";
+  throw std::invalid_argument(msg.str());
+}
+
 mx::array to_array(
     const ScalarOrArray& v,
     std::optional<mx::Dtype> dtype = std::nullopt);

python/tests/test_array.py

@@ -550,6 +550,29 @@ def test_construction_from_lists_wide_ints(self):
         # A float in the list still makes it float, not int64.
         self.assertEqual(mx.array([2**40, 1.5]).dtype, mx.float32)
 
+    def test_construction_from_python_ints_above_int64(self):
+        # Ints in (int64_max, uint64_max] widen to uint64, the same rule
+        # numpy uses. They used to raise std::bad_cast.
+        for value in (2**63, 2**64 - 1):
+            a = mx.array(value)
+            self.assertEqual(a.dtype, mx.uint64, msg=str(value))
+            self.assertEqual(a.item(), value)
+            self.assertEqual(mx.array(value, mx.uint64).item(), value)
+            self.assertEqual(mx.full((2,), value).tolist(), [value, value])
+        self.assertEqual(
+            (mx.zeros((2,), dtype=mx.uint64) + 2**63).tolist(), [2**63, 2**63]
+        )
+        # Both int64 boundaries keep int64.
+        self.assertEqual(mx.array(2**63 - 1).dtype, mx.int64)
+        self.assertEqual(mx.array(-(2**63)).dtype, mx.int64)
+        # Values outside both ranges raise a clean error, not std::bad_cast.
+        for value in (2**64, -(2**63) - 1):
+            with self.assertRaises(ValueError):
+                mx.array(value)
+        # Lists do not widen to uint64, but the error is clean.
+        with self.assertRaises(ValueError):
+            mx.array([2**63])
+
     def test_construction_from_lists_of_mlx_arrays(self):
         dtypes = [
             mx.bool_,