zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

fix slice update indexing (#1053)

Browse Source
This commit is contained in:
Awni Hannun
2024-04-29 12:17:40 -07:00
committed by GitHub
parent 490c0c4fdc
commit 09f1777896
4 changed files with 38 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
mlx/ops.cpp
View File

@@ -426,7 +426,7 @@ array expand_dims(const array& a, int axis, StreamOrDevice s /* = {} */) {
int ax = axis < 0 ? axis + out_dim : axis;
if (ax < 0 || ax >= out_dim) {
std::ostringstream msg;
msg << "[expand_dims] Invalid axes " << axis << " for output array with "
msg << "[expand_dims] Invalid axis " << axis << " for output array with "
<< a.ndim() << " dimensions.";
throw std::invalid_argument(msg.str());
}
@@ -452,7 +452,7 @@ array expand_dims(
ax = ax < 0 ? ax + out_ndim : ax;
if (ax < 0 || ax >= out_ndim) {
std::ostringstream msg;
msg << "[expand_dims] Invalid axes " << ax << " for output array with "
msg << "[expand_dims] Invalid axis " << ax << " for output array with "
<< a.ndim() << " dimensions.";
throw std::invalid_argument(msg.str());
}
@@ -591,7 +591,6 @@ array slice_update(
if (!has_neg_strides && upd_shape == src.shape()) {
return astype(update_broadcasted, src.dtype(), s);
}
return array(
src.shape(),
src.dtype(),