fix jit reduce (#1395)

2025-12-16 01:49:05 +08:00 · 2024-09-04 14:03:10 -07:00
parent 969337345f
commit 41c603d48a
7 changed files with 88 additions and 75 deletions
--- a/mlx/backend/metal/nojit_kernels.cpp
+++ b/mlx/backend/metal/nojit_kernels.cpp
@@ -104,8 +104,12 @@ MTL::ComputePipelineState* get_reduce_kernel(
    metal::Device& d,
    const std::string& kernel_name,
    const std::string&,
+    const std::string&,
    const array&,
-    const array&) {
+    const array&,
+    int,
+    int,
+    int) {
  return d.get_kernel(kernel_name);
 }