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 row reduce with very few rows (#1447)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2024-09-29 20:00:35 -07:00
committed by GitHub
parent 4a64d4bff1
commit bacced53d3
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
mlx/backend/metal/reduce.cpp
View File

@@ -449,7 +449,8 @@ void row_reduce_general_dispatch(
}
// Case 2: Contiguous reduce without non-row reductions
if (plan.type == ContiguousReduce && args.reduce_ndim == 0) {
if (plan.type == ContiguousReduce && args.reduce_ndim == 0 &&
in.size() / args.row_size >= 32) {
return row_reduce_simple(in, out, op_name, args, compute_encoder, d, s);
}