zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-22 21:48:08 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Indexing bug fix (#947)

Browse Source 
* Fix axes accounting

* Add tests
This commit is contained in:
Jagrit Digani
2024-04-01 12:18:50 -07:00
committed by GitHub
parent 02fedbf1da
commit 639e06e1f3
2 changed files with 23 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/src/indexing.cpp
View File

@@ -606,8 +606,12 @@ std::tuple<std::vector<array>, array, std::vector<int>> mlx_scatter_args_nd(
auto slice = nb::cast<nb::slice>(idx);
int stride = get_slice_int(nb::getattr(slice, "step"), 1);
num_strided_slices += (stride != 1);
num_simple_slices_post += (stride == 1);
if (stride != 1) {
num_strided_slices++;
num_simple_slices_post = 0;
} else {
num_simple_slices_post++;
}
} else if (nb::isinstance<array>(idx)) {
have_array = true;