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

Fix slice data size (#1913)

Browse Source 
* fix slice data size

* add test
This commit is contained in:
Awni Hannun
2025-03-02 21:50:42 -08:00
committed by GitHub
parent 5e6c130d93
commit 4e7cd31d12
2 changed files with 13 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
mlx/backend/common/slicing.cpp
View File

@@ -14,6 +14,10 @@ std::tuple<int64_t, Strides> prepare_slice(
data_offset += start_indices[i] * in.strides()[i];
inp_strides[i] = in.strides()[i] * strides[i];
}
// Normalize the offset
if (data_offset < 0) {
data_offset += in.data_size();
}
return std::make_tuple(data_offset, inp_strides);
}
@@ -54,9 +58,10 @@ void slice(
data_end += end_idx * in.strides()[i];
}
}
// data_end can be -1
size_t data_size =
data_end < 0 ? (data_offset - data_end) : (data_end - data_offset);
if (data_end < 0) {
data_end += in.data_size();
}
size_t data_size = (data_end - data_offset);
shared_buffer_slice(in, inp_strides, data_offset, data_size, out);
}