Fix row reduce with very few rows (ml-explore#1447)

lighttransport · Sep 30, 2024 · bacced5 · bacced5
1 parent 4a64d4b
commit bacced5
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/mlx/backend/metal/reduce.cpp b/mlx/backend/metal/reduce.cpp
@@ -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);
   }