zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-24 20:28:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

2d gather specialization (#1339)

Browse Source
This commit is contained in:
Awni Hannun
2024-08-22 10:48:24 -07:00
committed by GitHub
parent 82db84b899
commit df3233454d
3 changed files with 34 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
mlx/backend/metal/jit/indexing.h
View File

@@ -13,8 +13,8 @@ constexpr std::string_view gather_kernels = R"(
const constant size_t* idx_strides [[buffer(8)]],
const constant int& idx_ndim [[buffer(9)]],
{4}
uint2 index [[thread_position_in_grid]],
uint2 grid_dim [[threads_per_grid]]) {{
uint3 index [[thread_position_in_grid]],
uint3 grid_dim [[threads_per_grid]]) {{
Indices<{2}, {3}> idxs{{
{{ {5} }}, idx_shapes, idx_strides, idx_ndim}};