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);