zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-26 02:33:21 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix upon review

Browse Source 
Update python/src/ops.cpp

Co-authored-by: Awni Hannun <awni.hannun@gmail.com>
This commit is contained in:
Hyunsung Lee 2025-04-22 05:47:07 +09:00
parent 3787167fb2
commit 5e8712cf6f
1 changed files with 5 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
python/src/ops.cpp
View File

@ -5194,26 +5194,18 @@ void init_ops(nb::module_& m) {
[](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] Must provide at least one shape.");
std::vector<mx::Shape> shape_vec;
shape_vec.reserve(shapes.size());
mx::Shape result;
for (size_t i = 0; i < shapes.size(); ++i) { for (size_t i = 0; i < shapes.size(); ++i) {
if (!nb::isinstance<mx::Shape>(shapes[i]) && if (!nb::isinstance<mx::Shape>(shapes[i]) &&
!nb::isinstance<nb::tuple>(shapes[i])) !nb::isinstance<nb::tuple>(shapes[i]))
throw std::invalid_argument( throw std::invalid_argument(
"broadcast_shapes expects a sequence of shapes (tuple or list of ints)"); "[broadcast_shapes] Expects a sequence of shapes (tuple or list of ints).");
result = mx::broadcast_shapes(result, nb::cast<mx::Shape>(shapes[i]));
shape_vec.push_back(nb::cast<mx::Shape>(shapes[i]));
} }
mx::Shape result = shape_vec[0]; return nb::tuple(nb::cast(result));
for (size_t i = 1; i < shape_vec.size(); ++i)
result = mx::broadcast_shapes(result, shape_vec[i]);
auto py_list = nb::cast(result);
return nb::tuple(py_list);
}, },
nb::sig("def broadcast_shapes(*shapes: Sequence[int]) -> Tuple[int]"), nb::sig("def broadcast_shapes(*shapes: Sequence[int]) -> Tuple[int]"),
R"pbdoc( R"pbdoc(