zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-12 19:11:19 +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());
for (size_t i = 0; i < shapes.size(); ++i) {
mx::Shape shape;
if (nb::isinstance<mx::Shape>(shapes[i]) ||
nb::isinstance<nb::tuple>(shapes[i])) {
shape = nb::cast<mx::Shape>(shapes[i]);
} else {
if (!nb::isinstance<mx::Shape>(shapes[i]) &&
!nb::isinstance<nb::tuple>(shapes[i]))
throw std::invalid_argument(
"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];