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/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,