fix bug with multiple attributes (#1348)

Co-authored-by: Alex Barron <[email protected]>
ml-explore · Aug 23, 2024 · da8deb2 · da8deb2
1 parent 98b6ce3
commit da8deb2
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/mlx/fast.cpp b/mlx/fast.cpp
@@ -1052,13 +1052,15 @@ void write_signature(
     index++;
   }
   // Add metal attributes e.g. `threadgroup_index_in_grid`
+  index = 0;
   for (const auto& [attr, dtype] : attrs) {
     kernel_source << "  " << dtype << " " << attr << " [[" << attr << "]]";
     if (index < attrs.size() - 1) {
       kernel_source << "," << std::endl;
     } else {
       kernel_source << ") {" << std::endl;
     }
+    index++;
   }
   kernel_source << source << std::endl;
   kernel_source << "}" << std::endl;

diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py
@@ -618,12 +618,12 @@ def test_custom_kernel_strides(self):
             uint elem = thread_position_in_grid.x;
             uint loc = elem_to_loc(elem, inp_shape, inp_strides, inp_ndim);
             T tmp = inp[loc];
-            out[elem] = metal::exp(tmp);
+            out[elem] = metal::exp(tmp) * threads_per_simdgroup;
         """
         source_contig = """
             uint elem = thread_position_in_grid.x;
             T tmp = inp[elem];
-            out[elem] = metal::exp(tmp);
+            out[elem] = metal::exp(tmp) * threads_per_simdgroup;
         """
 
         # non contiguous
@@ -644,7 +644,7 @@ def test_custom_kernel_strides(self):
                 output_dtypes={"out": a.dtype},
                 stream=mx.gpu,
             )
-            self.assertTrue(mx.allclose(mx.exp(a), outputs["out"]))
+            self.assertTrue(mx.allclose(mx.exp(a) * 32, outputs["out"]))
 
 
 if __name__ == "__main__":