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

nit

Browse Source
This commit is contained in:
Hyunsung Lee 2025-04-20 22:47:37 +09:00
parent ed777f2ec6
commit 3787167fb2
1 changed files with 3 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
python/src/ops.cpp
View File

@ -5200,16 +5200,12 @@ void init_ops(nb::module_& m) {
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; 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])) {
shape = nb::cast<mx::Shape>(shapes[i]);
} else {
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)");
}
shape_vec.push_back(shape); shape_vec.push_back(nb::cast<mx::Shape>(shapes[i]));
} }
mx::Shape result = shape_vec[0]; mx::Shape result = shape_vec[0];