L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Clearer error when shape dimension overflows int32 (#3425)

Browse Source 
Co-authored-by: Kanishk <kanishk.chores@gmail.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
serenposh
2026-05-05 06:23:36 +05:30
committed by GitHub
parent b43965925f
commit 1fdd4e23c2
6 changed files with 122 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/src/convert.cpp
+12 -3
View File
@@ -1,5 +1,8 @@
// Copyright © 2024 Apple Inc.
#include <limits>
#include <sstream>
#include <nanobind/stl/complex.h>
#include "python/src/convert.h"
@@ -15,9 +18,15 @@ enum PyScalarT {
};
int check_shape_dim(int64_t dim) {
if (dim > std::numeric_limits<int>::max()) {
throw std::invalid_argument(
"Shape dimension falls outside supported `int` range.");
if (dim > std::numeric_limits<int>::max() ||
dim < std::numeric_limits<int>::min()) {
std::ostringstream msg;
msg << "Shape dimension " << dim << " is outside the supported range ["
<< std::numeric_limits<int>::min() << ", "
<< std::numeric_limits<int>::max()
<< "]. MLX currently uses 32-bit integers for shape dimensions.";
PyErr_SetString(PyExc_OverflowError, msg.str().c_str());
nb::detail::raise_python_error();
}
return static_cast<int>(dim);
}
python/src/convert.h
+4
View File
@@ -76,3 +76,7 @@ nb::object tolist(mx::array& a);
mx::array create_array(nb::object v, std::optional<mx::Dtype> t);
mx::array array_from_list(nb::list pl, std::optional<mx::Dtype> dtype);
mx::array array_from_list(nb::tuple pl, std::optional<mx::Dtype> dtype);
// Narrow a Python-side shape dimension (int64) to a C++ mx::ShapeElem (int32),
// raising a clear error if the value would overflow.
int check_shape_dim(int64_t dim);
python/src/ops.cpp
+14 -18
View File
@@ -15,6 +15,7 @@
#include "mlx/einsum.h"
#include "mlx/ops.h"
#include "mlx/utils.h"
#include "python/src/convert.h"
#include "python/src/load.h"
#include "python/src/small_vector.h"
#include "python/src/utils.h"
@@ -45,6 +46,13 @@ double scalar_to_double(Scalar s) {
}
}
mx::Shape to_shape(const nb::object& shape) {
if (nb::isinstance<nb::int_>(shape)) {
return {check_shape_dim(nb::cast<int64_t>(shape))};
}
return nb::cast<mx::Shape>(shape);
}
void init_ops(nb::module_& m) {
m.def(
"reshape",
@@ -1702,15 +1710,11 @@ void init_ops(nb::module_& m) {
)pbdoc");
m.def(
"full",
[](const std::variant<int, mx::Shape>& shape,
[](const nb::object& shape,
const ScalarOrArray& vals,
std::optional<mx::Dtype> dtype,
mx::StreamOrDevice s) {
if (auto pv = std::get_if<int>(&shape); pv) {
return mx::full({*pv}, to_array(vals, dtype), s);
} else {
return mx::full(std::get<mx::Shape>(shape), to_array(vals, dtype), s);
}
return mx::full(to_shape(shape), to_array(vals, dtype), s);
},
"shape"_a,
"vals"_a,
@@ -1736,15 +1740,11 @@ void init_ops(nb::module_& m) {
)pbdoc");
m.def(
"zeros",
[](const std::variant<int, mx::Shape>& shape,
[](const nb::object& shape,
std::optional<mx::Dtype> dtype,
mx::StreamOrDevice s) {
auto t = dtype.value_or(mx::float32);
if (auto pv = std::get_if<int>(&shape); pv) {
return mx::zeros({*pv}, t, s);
} else {
return mx::zeros(std::get<mx::Shape>(shape), t, s);
}
return mx::zeros(to_shape(shape), t, s);
},
"shape"_a,
"dtype"_a.none() = mx::float32,
@@ -1802,15 +1802,11 @@ void init_ops(nb::module_& m) {
)pbdoc");
m.def(
"ones",
[](const std::variant<int, mx::Shape>& shape,
[](const nb::object& shape,
std::optional<mx::Dtype> dtype,
mx::StreamOrDevice s) {
auto t = dtype.value_or(mx::float32);
if (auto pv = std::get_if<int>(&shape); pv) {
return mx::ones({*pv}, t, s);
} else {
return mx::ones(std::get<mx::Shape>(shape), t, s);
}
return mx::ones(to_shape(shape), t, s);
},
"shape"_a,
"dtype"_a.none() = mx::float32,
python/src/small_vector.h
+41 -8
View File
@@ -2,6 +2,11 @@
#pragma once
#include <cstdint>
#include <limits>
#include <sstream>
#include <type_traits>
#include "mlx/small_vector.h"
#include <nanobind/stl/detail/nb_list.h>
@@ -14,11 +19,19 @@ struct type_caster<mlx::core::SmallVector<Type, Size, Alloc>> {
using List = mlx::core::SmallVector<Type, Size, Alloc>;
using Caster = make_caster<Type>;
// For narrow integer element types we fetch each element through a wider
// integer caster so we can emit a clean OverflowError on overflow instead of
// nanobind's generic "incompatible function arguments" TypeError.
static constexpr bool kNarrowInt = std::is_integral_v<Type> &&
!std::is_same_v<Type, bool> && (sizeof(Type) < sizeof(int64_t));
NB_TYPE_CASTER(
List,
const_name("tuple[") + make_caster<Type>::Name + const_name(", ...]"))
bool from_python(handle src, uint8_t flags, cleanup_list* cleanup) noexcept {
// Not noexcept: on overflow of a narrow integer element we raise
// OverflowError so nanobind surfaces a clean error to the user.
bool from_python(handle src, uint8_t flags, cleanup_list* cleanup) {
size_t size;
PyObject* temp;
@@ -29,19 +42,39 @@ struct type_caster<mlx::core::SmallVector<Type, Size, Alloc>> {
value.clear();
value.reserve(size);
Caster caster;
bool success = o != nullptr;
flags = flags_for_local_caster<Type>(flags);
for (size_t i = 0; i < size; ++i) {
if (!caster.from_python(o[i], flags, cleanup) ||
!caster.template can_cast<Type>()) {
success = false;
break;
if constexpr (kNarrowInt) {
make_caster<int64_t> wide;
if (!wide.from_python(o[i], flags, cleanup) ||
!wide.template can_cast<int64_t>()) {
success = false;
break;
}
int64_t v = wide.operator cast_t<int64_t>();
if (v > std::numeric_limits<Type>::max() ||
v < std::numeric_limits<Type>::min()) {
std::ostringstream msg;
msg << "Integer value " << v << " is outside the supported range ["
<< static_cast<int64_t>(std::numeric_limits<Type>::min()) << ", "
<< static_cast<int64_t>(std::numeric_limits<Type>::max()) << "].";
Py_XDECREF(temp);
PyErr_SetString(PyExc_OverflowError, msg.str().c_str());
raise_python_error();
}
value.push_back(static_cast<Type>(v));
} else {
Caster caster;
if (!caster.from_python(o[i], flags, cleanup) ||
!caster.template can_cast<Type>()) {
success = false;
break;
}
value.push_back(caster.operator cast_t<Type>());
}
value.push_back(caster.operator cast_t<Type>());
}
Py_XDECREF(temp);
python/tests/test_array.py
+5 -2
View File
@@ -767,10 +767,13 @@ class TestArray(mlx_tests.MLXTestCase):
def test_array_np_shape_dim_check(self):
a_npy = np.empty(2**31, dtype=np.bool_)
with self.assertRaises(ValueError) as e:
with self.assertRaises(OverflowError) as e:
mx.array(a_npy)
self.assertEqual(
str(e.exception), "Shape dimension falls outside supported `int` range."
str(e.exception),
"Shape dimension 2147483648 is outside the supported range "
"[-2147483648, 2147483647]. MLX currently uses 32-bit integers "
"for shape dimensions.",
)
def test_dtype_promotion(self):
python/tests/test_ops.py
+46
View File
@@ -92,6 +92,52 @@ class TestOps(mlx_tests.MLXTestCase):
self.assertEqual(y.dtype, t)
self.assertTrue(mx.array_equal(y, x))
def test_shape_overflow_error(self):
# Shape dimensions that don't fit in int32 should raise a clear
# OverflowError that names the offending value, rather than a generic
# "incompatible function arguments" TypeError. The overflow check
# lives in the mx::Shape type caster, so it applies to every op that
# takes a shape. See issue #2681.
too_big = 2**31
# Array creation ops — also exercise the scalar shape path.
for ctor in (mx.zeros, mx.ones):
with self.assertRaises(OverflowError) as cm:
ctor(too_big)
self.assertIn(str(too_big), str(cm.exception))
with self.assertRaises(OverflowError) as cm:
ctor([too_big])
self.assertIn(str(too_big), str(cm.exception))
with self.assertRaises(OverflowError) as cm:
mx.full(too_big, 0.0)
self.assertIn(str(too_big), str(cm.exception))
with self.assertRaises(OverflowError) as cm:
mx.full([too_big], 0.0)
self.assertIn(str(too_big), str(cm.exception))
# Other shape-taking ops should surface the same clean error.
a = mx.zeros(4)
with self.assertRaises(OverflowError) as cm:
mx.reshape(a, [too_big])
self.assertIn(str(too_big), str(cm.exception))
with self.assertRaises(OverflowError) as cm:
mx.broadcast_to(a, [too_big, 1])
self.assertIn(str(too_big), str(cm.exception))
# Negative overflow (< int32 min) is caught too.
too_negative = -(2**31) - 1
with self.assertRaises(OverflowError) as cm:
mx.zeros([too_negative])
self.assertIn(str(too_negative), str(cm.exception))
# Shapes that fit in int32 still go through unchanged.
self.assertEqual(mx.zeros(4).shape, (4,))
self.assertEqual(mx.zeros((2, 3)).shape, (2, 3))
self.assertEqual(mx.ones([2, 3]).shape, (2, 3))
self.assertEqual(mx.full((2, 3), 1.5).tolist(), [[1.5] * 3] * 2)
def test_scalar_inputs(self):
# Check combinations of python types
a = mx.add(False, True)