zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-24 02:41:19 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

nit

Browse Source
This commit is contained in:
Hyunsung Lee 2025-04-20 12:58:26 +09:00
parent a7a96b0ad6
commit 92444f393b
1 changed files with 2 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
python/src/ops.cpp
View File

@ -5192,17 +5192,15 @@ void init_ops(nb::module_& m) {
m.def( m.def(
"broadcast_shapes", "broadcast_shapes",
[](const nb::args& shapes) { [](const nb::args& shapes) {
if (shapes.size() == 0) { if (shapes.size() == 0)
throw std::invalid_argument( throw std::invalid_argument(
"broadcast_shapes expects a sequence of shapes"); "broadcast_shapes expects a sequence of shapes");
}
std::vector<mx::Shape> shape_vec; std::vector<mx::Shape> shape_vec;
shape_vec.reserve(shapes.size()); shape_vec.reserve(shapes.size());
for (size_t i = 0; i < shapes.size(); ++i) { for (size_t i = 0; i < shapes.size(); ++i) {
mx::Shape shape; mx::Shape shape;
if (nb::isinstance<nb::tuple>(shapes[i])) { if (nb::isinstance<nb::tuple>(shapes[i])) {
nb::tuple t = nb::cast<nb::tuple>(shapes[i]); nb::tuple t = nb::cast<nb::tuple>(shapes[i]);
for (size_t j = 0; j < t.size(); ++j) { for (size_t j = 0; j < t.size(); ++j) {
@ -5221,14 +5219,9 @@ void init_ops(nb::module_& m) {
shape_vec.push_back(shape); shape_vec.push_back(shape);
} }
if (shape_vec.empty()) {
return nb::tuple();
}
mx::Shape result = shape_vec[0]; mx::Shape result = shape_vec[0];
for (size_t i = 1; i < shape_vec.size(); ++i) { for (size_t i = 1; i < shape_vec.size(); ++i)
result = mx::broadcast_shapes(result, shape_vec[i]); result = mx::broadcast_shapes(result, shape_vec[i]);
}
auto py_list = nb::cast(result); auto py_list = nb::cast(result);
return nb::tuple(py_list); return nb::tuple(py_list);