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 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
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 1 deletions
Show Stats 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);
}