diff --git a/python/src/convert.cpp b/python/src/convert.cpp index edc40ff8..471ff24a 100644 --- a/python/src/convert.cpp +++ b/python/src/convert.cpp @@ -1,5 +1,8 @@ // Copyright © 2024 Apple Inc. +#include +#include + #include #include "python/src/convert.h" @@ -15,9 +18,15 @@ enum PyScalarT { }; int check_shape_dim(int64_t dim) { - if (dim > std::numeric_limits::max()) { - throw std::invalid_argument( - "Shape dimension falls outside supported `int` range."); + if (dim > std::numeric_limits::max() || + dim < std::numeric_limits::min()) { + std::ostringstream msg; + msg << "Shape dimension " << dim << " is outside the supported range [" + << std::numeric_limits::min() << ", " + << std::numeric_limits::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(dim); } diff --git a/python/src/convert.h b/python/src/convert.h index 9da7a0f7..fa492d86 100644 --- a/python/src/convert.h +++ b/python/src/convert.h @@ -76,3 +76,7 @@ nb::object tolist(mx::array& a); mx::array create_array(nb::object v, std::optional t); mx::array array_from_list(nb::list pl, std::optional dtype); mx::array array_from_list(nb::tuple pl, std::optional 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); diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 9bbe0ff7..9a48b37a 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -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(shape)) { + return {check_shape_dim(nb::cast(shape))}; + } + return nb::cast(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& shape, + [](const nb::object& shape, const ScalarOrArray& vals, std::optional dtype, mx::StreamOrDevice s) { - if (auto pv = std::get_if(&shape); pv) { - return mx::full({*pv}, to_array(vals, dtype), s); - } else { - return mx::full(std::get(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& shape, + [](const nb::object& shape, std::optional dtype, mx::StreamOrDevice s) { auto t = dtype.value_or(mx::float32); - if (auto pv = std::get_if(&shape); pv) { - return mx::zeros({*pv}, t, s); - } else { - return mx::zeros(std::get(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& shape, + [](const nb::object& shape, std::optional dtype, mx::StreamOrDevice s) { auto t = dtype.value_or(mx::float32); - if (auto pv = std::get_if(&shape); pv) { - return mx::ones({*pv}, t, s); - } else { - return mx::ones(std::get(shape), t, s); - } + return mx::ones(to_shape(shape), t, s); }, "shape"_a, "dtype"_a.none() = mx::float32, diff --git a/python/src/small_vector.h b/python/src/small_vector.h index 2be8dcec..d868370e 100644 --- a/python/src/small_vector.h +++ b/python/src/small_vector.h @@ -2,6 +2,11 @@ #pragma once +#include +#include +#include +#include + #include "mlx/small_vector.h" #include @@ -14,11 +19,19 @@ struct type_caster> { using List = mlx::core::SmallVector; using Caster = make_caster; + // 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 && + !std::is_same_v && (sizeof(Type) < sizeof(int64_t)); + NB_TYPE_CASTER( List, const_name("tuple[") + make_caster::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> { value.clear(); value.reserve(size); - Caster caster; bool success = o != nullptr; flags = flags_for_local_caster(flags); for (size_t i = 0; i < size; ++i) { - if (!caster.from_python(o[i], flags, cleanup) || - !caster.template can_cast()) { - success = false; - break; + if constexpr (kNarrowInt) { + make_caster wide; + if (!wide.from_python(o[i], flags, cleanup) || + !wide.template can_cast()) { + success = false; + break; + } + int64_t v = wide.operator cast_t(); + if (v > std::numeric_limits::max() || + v < std::numeric_limits::min()) { + std::ostringstream msg; + msg << "Integer value " << v << " is outside the supported range [" + << static_cast(std::numeric_limits::min()) << ", " + << static_cast(std::numeric_limits::max()) << "]."; + Py_XDECREF(temp); + PyErr_SetString(PyExc_OverflowError, msg.str().c_str()); + raise_python_error(); + } + value.push_back(static_cast(v)); + } else { + Caster caster; + if (!caster.from_python(o[i], flags, cleanup) || + !caster.template can_cast()) { + success = false; + break; + } + value.push_back(caster.operator cast_t()); } - - value.push_back(caster.operator cast_t()); } Py_XDECREF(temp); diff --git a/python/tests/test_array.py b/python/tests/test_array.py index a35460c9..35d02a6d 100644 --- a/python/tests/test_array.py +++ b/python/tests/test_array.py @@ -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): diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index ccc26bdc..77db5303 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -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)