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

No extra reshape (#1557)

Browse Source 
* no extra reshape

* lint
This commit is contained in:
Awni Hannun 2024-11-02 19:07:20 -07:00 committed by GitHub
parent 46d8b16ab4
commit 09bc32f62f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
python/src/indexing.cpp
View File

@ -803,9 +803,10 @@ auto mlx_slice_update(
// Pre process tuple
auto upd = to_array(v, src.dtype());
// Remove leading singletons dimensions from the update
// Remove extra leading singletons dimensions from the update
int s = 0;
for (; s < upd.ndim() && upd.shape(s) == 1; s++) {
for (; s < upd.ndim() && upd.shape(s) == 1 && (upd.ndim() - s) > src.ndim();
s++) {
};
auto up_shape = std::vector<int>(upd.shape().begin() + s, upd.shape().end());
up_shape = up_shape.empty() ? std::vector{1} : up_shape;