L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Slice update with operation (#3266)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2026-03-18 06:18:02 -07:00
committed by GitHub
parent e353be8235
commit 7bc61cceed
18 changed files with 1475 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/src/indexing.cpp
+81 -31
View File
@@ -769,43 +769,53 @@ mlx_compute_scatter_args(
throw std::invalid_argument("Cannot index mlx array using the given type.");
}
auto mlx_slice_update(
std::tuple<std::optional<mx::array>, mx::Shape, mx::Shape, mx::Shape>
mlx_compute_slice_update_args(
const mx::array& src,
const nb::object& obj,
const ScalarOrArray& v) {
// Build the slice params
mx::Shape starts(src.ndim(), 0);
mx::Shape stops = src.shape();
mx::Shape strides(src.ndim(), 1);
// Can't route to slice update if not slice, tuple, or int
if (src.ndim() == 0 || nb::isinstance<nb::bool_>(obj) ||
(!nb::isinstance<nb::slice>(obj) && !nb::isinstance<nb::tuple>(obj) &&
!nb::isinstance<nb::int_>(obj))) {
return std::make_pair(false, src);
return std::make_tuple(
std::nullopt, std::move(starts), std::move(stops), std::move(strides));
}
if (nb::isinstance<nb::tuple>(obj)) {
// Can't route to slice update if any arrays are present
for (auto idx : nb::cast<nb::tuple>(obj)) {
if (nb::isinstance<mx::array>(idx) || nb::isinstance<nb::list>(idx)) {
return std::make_pair(false, src);
return std::make_tuple(
std::nullopt,
std::move(starts),
std::move(stops),
std::move(strides));
}
}
}
// Should be able to route to slice update
// Pre process tuple
auto upd = to_array(v, src.dtype());
// Should be able to route to slice update just extract the update value and
// and the slice arguments.
// Cast v to an array and ensure it is the right type
auto update = to_array(v, src.dtype());
// Remove extra leading singletons dimensions from the update
int s = 0;
for (; s < static_cast<int>(upd.ndim()) - 1 && upd.shape(s) == 1 &&
(upd.ndim() - s) > src.ndim();
for (; s < static_cast<int>(update.ndim()) - 1 && update.shape(s) == 1 &&
(update.ndim() - s) > src.ndim();
s++) {
};
auto squeeze_axes = std::vector<int>(s);
std::iota(squeeze_axes.begin(), squeeze_axes.end(), 0);
auto up = mx::squeeze(upd, squeeze_axes);
update = mx::squeeze(update, squeeze_axes);
// Build slice update params
mx::Shape starts(src.ndim(), 0);
mx::Shape stops = src.shape();
mx::Shape strides(src.ndim(), 1);
// Single int then make it a slice of size 1
if (nb::isinstance<nb::int_>(obj)) {
if (src.ndim() < 1) {
std::ostringstream msg;
@@ -816,12 +826,11 @@ auto mlx_slice_update(
idx = idx < 0 ? idx + stops[0] : idx;
starts[0] = idx;
stops[0] = idx + 1;
auto out = slice_update(
src, up, std::move(starts), std::move(stops), std::move(strides));
return std::make_pair(true, out);
return std::make_tuple(
update, std::move(starts), std::move(stops), std::move(strides));
}
// If it's just a simple slice, just do a slice update and return
// Simple slice, just extract it into the first dim
if (nb::isinstance<nb::slice>(obj)) {
// Read slice arguments
get_slice_params(
@@ -830,16 +839,14 @@ auto mlx_slice_update(
strides[0],
nb::cast<nb::slice>(obj),
src.shape(0));
// Do slice update
auto out = slice_update(src, up, starts, stops, strides);
return std::make_pair(true, out);
return std::make_tuple(
update, std::move(starts), std::move(stops), std::move(strides));
}
// It must be a tuple
auto entries = nb::cast<nb::tuple>(obj);
// Expand ellipses into a series of ':' slices
// Expand ellipsis into a series of ':' slices
auto [non_none_indices, indices] = mlx_expand_ellipsis(src.shape(), entries);
// Dimension check
@@ -851,15 +858,20 @@ auto mlx_slice_update(
// If no non-None indices return the broadcasted update
if (non_none_indices == 0) {
return std::make_pair(true, broadcast_to(up, src.shape()));
return std::make_tuple(
broadcast_to(update, src.shape()),
std::move(starts),
std::move(stops),
std::move(strides));
}
// Parse the update slice
int unspecified = src.ndim() - non_none_indices;
std::vector<int> squeeze_dims;
std::vector<int> expand_dims;
for (int i = indices.size() - 1,
ax = non_none_indices - 1,
upd_ax = upd.ndim() - unspecified - 1;
upd_ax = update.ndim() - unspecified - 1;
i >= 0;
--i) {
auto& pyidx = indices[i];
@@ -887,11 +899,11 @@ auto mlx_slice_update(
}
}
}
update = mx::squeeze(
mx::expand_dims(update, std::move(expand_dims)), std::move(squeeze_dims));
up = mx::squeeze(
mx::expand_dims(up, std::move(expand_dims)), std::move(squeeze_dims));
auto out = slice_update(src, up, starts, stops, strides);
return std::make_pair(true, out);
return std::make_tuple(
update, std::move(starts), std::move(stops), std::move(strides));
}
std::optional<mx::array> extract_boolean_mask(const nb::object& obj) {
@@ -921,9 +933,11 @@ void mlx_set_item(
mx::array& src,
const nb::object& obj,
const ScalarOrArray& v) {
auto [success, out] = mlx_slice_update(src, obj, v);
if (success) {
src.overwrite_descriptor(out);
auto [update, starts, stops, strides] =
mlx_compute_slice_update_args(src, obj, v);
if (update) {
src.overwrite_descriptor(
slice_update(src, *update, starts, stops, strides));
return;
}
@@ -947,6 +961,12 @@ mx::array mlx_add_item(
const mx::array& src,
const nb::object& obj,
const ScalarOrArray& v) {
auto [update, starts, stops, strides] =
mlx_compute_slice_update_args(src, obj, v);
if (update) {
return slice_update_add(src, *update, starts, stops, strides);
}
auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
if (indices.size() > 0) {
return scatter_add(src, indices, updates, axes);
@@ -959,6 +979,12 @@ mx::array mlx_subtract_item(
const mx::array& src,
const nb::object& obj,
const ScalarOrArray& v) {
auto [update, starts, stops, strides] =
mlx_compute_slice_update_args(src, obj, v);
if (update) {
return slice_update_add(src, -(*update), starts, stops, strides);
}
auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
if (indices.size() > 0) {
return scatter_add(src, indices, -updates, axes);
@@ -971,6 +997,12 @@ mx::array mlx_multiply_item(
const mx::array& src,
const nb::object& obj,
const ScalarOrArray& v) {
auto [update, starts, stops, strides] =
mlx_compute_slice_update_args(src, obj, v);
if (update) {
return slice_update_prod(src, *update, starts, stops, strides);
}
auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
if (indices.size() > 0) {
return scatter_prod(src, indices, updates, axes);
@@ -983,6 +1015,12 @@ mx::array mlx_divide_item(
const mx::array& src,
const nb::object& obj,
const ScalarOrArray& v) {
auto [update, starts, stops, strides] =
mlx_compute_slice_update_args(src, obj, v);
if (update) {
return slice_update_prod(src, reciprocal(*update), starts, stops, strides);
}
auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
if (indices.size() > 0) {
return scatter_prod(src, indices, reciprocal(updates), axes);
@@ -995,6 +1033,12 @@ mx::array mlx_maximum_item(
const mx::array& src,
const nb::object& obj,
const ScalarOrArray& v) {
auto [update, starts, stops, strides] =
mlx_compute_slice_update_args(src, obj, v);
if (update) {
return slice_update_max(src, *update, starts, stops, strides);
}
auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
if (indices.size() > 0) {
return scatter_max(src, indices, updates, axes);
@@ -1007,6 +1051,12 @@ mx::array mlx_minimum_item(
const mx::array& src,
const nb::object& obj,
const ScalarOrArray& v) {
auto [update, starts, stops, strides] =
mlx_compute_slice_update_args(src, obj, v);
if (update) {
return slice_update_min(src, *update, starts, stops, strides);
}
auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
if (indices.size() > 0) {
return scatter_min(src, indices, updates, axes);