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

Fix binary op dispatch (#584)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2024-01-29 19:36:17 -08:00
committed by GitHub
parent 3c2f192345
commit 65d0b8df9f
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
mlx/backend/metal/primitives.cpp
View File

@@ -60,7 +60,7 @@ void binary_op(
break;
}
kname << op << type_to_name(a);
if (bopt == General && out.ndim() <= MAX_BINARY_SPECIALIZED_DIMS) {
if (bopt == General && shape.size() <= MAX_BINARY_SPECIALIZED_DIMS) {
kname << "_" << shape.size();
}
@@ -158,7 +158,7 @@ void binary_op(
break;
}
kname << op << type_to_name(a);
if (bopt == General && out.ndim() <= MAX_BINARY_SPECIALIZED_DIMS) {
if (bopt == General && shape.size() <= MAX_BINARY_SPECIALIZED_DIMS) {
kname << "_" << shape.size();
}