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

Remove "using namespace mlx::core" in python/src (#1689)

Browse Source
This commit is contained in:
Cheng
2024-12-12 08:45:39 +09:00
committed by GitHub
parent f3dfa36a3a
commit 0bf19037ca
22 changed files with 1423 additions and 1302 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/src/indexing.cpp
+73 -66
View File
@@ -43,20 +43,20 @@ void get_slice_params(
nb::getattr(in_slice, "stop"), strides < 0 ? -axis_size - 1 : axis_size);
}
array get_int_index(nb::object idx, int axis_size) {
mx::array get_int_index(nb::object idx, int axis_size) {
int idx_ = nb::cast<int>(idx);
idx_ = (idx_ < 0) ? idx_ + axis_size : idx_;
return array(idx_, uint32);
return mx::array(idx_, mx::uint32);
}
bool is_valid_index_type(const nb::object& obj) {
return nb::isinstance<nb::slice>(obj) || nb::isinstance<nb::int_>(obj) ||
nb::isinstance<array>(obj) || obj.is_none() || nb::ellipsis().is(obj) ||
nb::isinstance<nb::list>(obj);
nb::isinstance<mx::array>(obj) || obj.is_none() ||
nb::ellipsis().is(obj) || nb::isinstance<nb::list>(obj);
}
array mlx_get_item_slice(const array& src, const nb::slice& in_slice) {
mx::array mlx_get_item_slice(const mx::array& src, const nb::slice& in_slice) {
// Check input and raise error if 0 dim for parity with np
if (src.ndim() == 0) {
throw std::invalid_argument(
@@ -77,14 +77,14 @@ array mlx_get_item_slice(const array& src, const nb::slice& in_slice) {
return slice(src, starts, ends, strides);
}
array mlx_get_item_array(const array& src, const array& indices) {
mx::array mlx_get_item_array(const mx::array& src, const mx::array& indices) {
// Check input and raise error if 0 dim for parity with np
if (src.ndim() == 0) {
throw std::invalid_argument(
"too many indices for array: array is 0-dimensional");
}
if (indices.dtype() == bool_) {
if (indices.dtype() == mx::bool_) {
throw std::invalid_argument("boolean indices are not yet supported");
}
@@ -93,7 +93,7 @@ array mlx_get_item_array(const array& src, const array& indices) {
return take(src, indices, 0);
}
array mlx_get_item_int(const array& src, const nb::int_& idx) {
mx::array mlx_get_item_int(const mx::array& src, const nb::int_& idx) {
// Check input and raise error if 0 dim for parity with np
if (src.ndim() == 0) {
throw std::invalid_argument(
@@ -105,13 +105,13 @@ array mlx_get_item_int(const array& src, const nb::int_& idx) {
return take(src, get_int_index(idx, src.shape(0)), 0);
}
array mlx_gather_nd(
array src,
mx::array mlx_gather_nd(
mx::array src,
const std::vector<nb::object>& indices,
bool gather_first,
int& max_dims) {
max_dims = 0;
std::vector<array> gather_indices;
std::vector<mx::array> gather_indices;
std::vector<bool> is_slice(indices.size(), false);
int num_slices = 0;
// gather all the arrays
@@ -127,13 +127,13 @@ array mlx_gather_nd(
start = (start < 0) ? start + src.shape(i) : start;
end = (end < 0) ? end + src.shape(i) : end;
gather_indices.push_back(arange(start, end, stride, uint32));
gather_indices.push_back(arange(start, end, stride, mx::uint32));
num_slices++;
is_slice[i] = true;
} else if (nb::isinstance<nb::int_>(idx)) {
gather_indices.push_back(get_int_index(idx, src.shape(i)));
} else if (nb::isinstance<array>(idx)) {
auto arr = nb::cast<array>(idx);
} else if (nb::isinstance<mx::array>(idx)) {
auto arr = nb::cast<mx::array>(idx);
max_dims = std::max(static_cast<int>(arr.ndim()), max_dims);
gather_indices.push_back(arr);
}
@@ -144,7 +144,7 @@ array mlx_gather_nd(
int slice_index = 0;
for (int i = 0; i < gather_indices.size(); i++) {
if (is_slice[i]) {
Shape index_shape(max_dims + num_slices, 1);
mx::Shape index_shape(max_dims + num_slices, 1);
index_shape[max_dims + slice_index] = gather_indices[i].shape(0);
gather_indices[i] = reshape(gather_indices[i], std::move(index_shape));
slice_index++;
@@ -158,7 +158,7 @@ array mlx_gather_nd(
// reshape them so that the int/array indices are last
for (int i = 0; i < gather_indices.size(); i++) {
if (i < num_slices) {
Shape index_shape(max_dims + num_slices, 1);
mx::Shape index_shape(max_dims + num_slices, 1);
index_shape[i] = gather_indices[i].shape(0);
gather_indices[i] = reshape(gather_indices[i], std::move(index_shape));
}
@@ -241,7 +241,7 @@ auto mlx_expand_ellipsis(
return std::make_pair(non_none_indices, indices);
}
array mlx_get_item_nd(array src, const nb::tuple& entries) {
mx::array mlx_get_item_nd(mx::array src, const nb::tuple& entries) {
// No indices make this a noop
if (entries.size() == 0) {
return src;
@@ -281,7 +281,7 @@ array mlx_get_item_nd(array src, const nb::tuple& entries) {
bool have_non_array = false;
bool gather_first = false;
for (auto& idx : indices) {
if (nb::isinstance<array>(idx) || (nb::isinstance<nb::int_>(idx))) {
if (nb::isinstance<mx::array>(idx) || (nb::isinstance<nb::int_>(idx))) {
if (have_array && have_non_array) {
gather_first = true;
break;
@@ -294,7 +294,7 @@ array mlx_get_item_nd(array src, const nb::tuple& entries) {
int n_arr = 0;
for (auto& idx : indices) {
n_arr += nb::isinstance<array>(idx);
n_arr += nb::isinstance<mx::array>(idx);
}
have_array &= n_arr > 0;
@@ -304,7 +304,7 @@ array mlx_get_item_nd(array src, const nb::tuple& entries) {
// Then find the last array
for (last_array = indices.size() - 1; last_array >= 0; last_array--) {
auto& idx = indices[last_array];
if (nb::isinstance<array>(idx) || nb::isinstance<nb::int_>(idx)) {
if (nb::isinstance<mx::array>(idx) || nb::isinstance<nb::int_>(idx)) {
break;
}
}
@@ -340,7 +340,7 @@ array mlx_get_item_nd(array src, const nb::tuple& entries) {
} else {
for (int i = 0; i < indices.size(); i++) {
auto& idx = indices[i];
if (nb::isinstance<array>(idx) || nb::isinstance<nb::int_>(idx)) {
if (nb::isinstance<mx::array>(idx) || nb::isinstance<nb::int_>(idx)) {
break;
} else if (idx.is_none()) {
remaining_indices.push_back(idx);
@@ -426,11 +426,11 @@ array mlx_get_item_nd(array src, const nb::tuple& entries) {
return src;
}
array mlx_get_item(const array& src, const nb::object& obj) {
mx::array mlx_get_item(const mx::array& src, const nb::object& obj) {
if (nb::isinstance<nb::slice>(obj)) {
return mlx_get_item_slice(src, nb::cast<nb::slice>(obj));
} else if (nb::isinstance<array>(obj)) {
return mlx_get_item_array(src, nb::cast<array>(obj));
} else if (nb::isinstance<mx::array>(obj)) {
return mlx_get_item_array(src, nb::cast<mx::array>(obj));
} else if (nb::isinstance<nb::int_>(obj)) {
return mlx_get_item_int(src, nb::cast<nb::int_>(obj));
} else if (nb::isinstance<nb::tuple>(obj)) {
@@ -448,10 +448,11 @@ array mlx_get_item(const array& src, const nb::object& obj) {
throw std::invalid_argument("Cannot index mlx array using the given type.");
}
std::tuple<std::vector<array>, array, std::vector<int>> mlx_scatter_args_int(
const array& src,
std::tuple<std::vector<mx::array>, mx::array, std::vector<int>>
mlx_scatter_args_int(
const mx::array& src,
const nb::int_& idx,
const array& update) {
const mx::array& update) {
if (src.ndim() == 0) {
throw std::invalid_argument(
"too many indices for array: array is 0-dimensional");
@@ -473,10 +474,11 @@ std::tuple<std::vector<array>, array, std::vector<int>> mlx_scatter_args_int(
{0}};
}
std::tuple<std::vector<array>, array, std::vector<int>> mlx_scatter_args_array(
const array& src,
const array& indices,
const array& update) {
std::tuple<std::vector<mx::array>, mx::array, std::vector<int>>
mlx_scatter_args_array(
const mx::array& src,
const mx::array& indices,
const mx::array& update) {
if (src.ndim() == 0) {
throw std::invalid_argument(
"too many indices for array: array is 0-dimensional");
@@ -500,10 +502,11 @@ std::tuple<std::vector<array>, array, std::vector<int>> mlx_scatter_args_array(
return {{indices}, up, {0}};
}
std::tuple<std::vector<array>, array, std::vector<int>> mlx_scatter_args_slice(
const array& src,
std::tuple<std::vector<mx::array>, mx::array, std::vector<int>>
mlx_scatter_args_slice(
const mx::array& src,
const nb::slice& in_slice,
const array& update) {
const mx::array& update) {
// Check input and raise error if 0 dim for parity with np
if (src.ndim() == 0) {
throw std::invalid_argument(
@@ -539,7 +542,7 @@ std::tuple<std::vector<array>, array, std::vector<int>> mlx_scatter_args_slice(
auto up = reshape(update, up_shape);
// Build array to mark start of slice
auto idx = array({start}, {1}, uint32);
auto idx = mx::array({start}, {1}, mx::uint32);
// Get slice size
int slice_size = (end - start);
@@ -551,20 +554,21 @@ std::tuple<std::vector<array>, array, std::vector<int>> mlx_scatter_args_slice(
up = broadcast_to(up, up_shape_broadcast);
auto indices = std::vector<array>{idx};
auto indices = std::vector<mx::array>{idx};
auto axes = std::vector<int>{0};
return {indices, up, axes};
}
return mlx_scatter_args_array(
src, arange(start, end, stride, uint32), update);
src, arange(start, end, stride, mx::uint32), update);
}
std::tuple<std::vector<array>, array, std::vector<int>> mlx_scatter_args_nd(
const array& src,
std::tuple<std::vector<mx::array>, mx::array, std::vector<int>>
mlx_scatter_args_nd(
const mx::array& src,
const nb::tuple& entries,
const array& update) {
const mx::array& update) {
// Expand ellipses into a series of ':' slices
auto [non_none_indices, indices] = mlx_expand_ellipsis(src.shape(), entries);
@@ -623,12 +627,12 @@ std::tuple<std::vector<array>, array, std::vector<int>> mlx_scatter_args_nd(
num_simple_slices_post++;
}
} else if (nb::isinstance<array>(idx)) {
} else if (nb::isinstance<mx::array>(idx)) {
have_array = true;
if (have_array && have_non_array) {
arrays_first = true;
}
max_dim = std::max(nb::cast<array>(idx).ndim(), max_dim);
max_dim = std::max(nb::cast<mx::array>(idx).ndim(), max_dim);
num_arrays++;
num_simple_slices_post = 0;
}
@@ -643,7 +647,7 @@ std::tuple<std::vector<array>, array, std::vector<int>> mlx_scatter_args_nd(
idx_ndim = idx_ndim == 0 ? 1 : idx_ndim;
// Go over each index type and translate to the needed scatter args
std::vector<array> arr_indices;
std::vector<mx::array> arr_indices;
int slice_num = 0;
int array_num = 0;
int ax = 0;
@@ -668,7 +672,7 @@ std::tuple<std::vector<array>, array, std::vector<int>> mlx_scatter_args_nd(
// If it's a simple slice, we only need to add the start index
if (array_num >= num_arrays && num_strided_slices <= 0 && stride == 1) {
auto idx = array({start}, idx_shape, uint32);
auto idx = mx::array({start}, idx_shape, mx::uint32);
slice_shapes.push_back(end - start);
arr_indices.push_back(idx);
@@ -677,7 +681,7 @@ std::tuple<std::vector<array>, array, std::vector<int>> mlx_scatter_args_nd(
}
// Otherwise we expand the slice into indices using arange
else {
auto idx = arange(start, end, stride, uint32);
auto idx = arange(start, end, stride, mx::uint32);
auto loc = slice_num + (arrays_first ? max_dim : 0);
idx_shape[loc] = idx.size();
arr_indices.push_back(reshape(idx, idx_shape));
@@ -696,9 +700,9 @@ std::tuple<std::vector<array>, array, std::vector<int>> mlx_scatter_args_nd(
} else if (pyidx.is_none()) {
// We only use the None's for bookeeping dimensions
slice_num++;
} else if (nb::isinstance<array>(pyidx)) {
} else if (nb::isinstance<mx::array>(pyidx)) {
ax++;
auto idx = nb::cast<array>(pyidx);
auto idx = nb::cast<mx::array>(pyidx);
std::vector<int> idx_shape(idx_ndim, 1);
// Place the arrays in the correct dimension
@@ -748,16 +752,16 @@ std::tuple<std::vector<array>, array, std::vector<int>> mlx_scatter_args_nd(
return {arr_indices, up, axes};
}
std::tuple<std::vector<array>, array, std::vector<int>>
std::tuple<std::vector<mx::array>, mx::array, std::vector<int>>
mlx_compute_scatter_args(
const array& src,
const mx::array& src,
const nb::object& obj,
const ScalarOrArray& v) {
auto vals = to_array(v, src.dtype());
if (nb::isinstance<nb::slice>(obj)) {
return mlx_scatter_args_slice(src, nb::cast<nb::slice>(obj), vals);
} else if (nb::isinstance<array>(obj)) {
return mlx_scatter_args_array(src, nb::cast<array>(obj), vals);
} else if (nb::isinstance<mx::array>(obj)) {
return mlx_scatter_args_array(src, nb::cast<mx::array>(obj), vals);
} else if (nb::isinstance<nb::int_>(obj)) {
return mlx_scatter_args_int(src, nb::cast<nb::int_>(obj), vals);
} else if (nb::isinstance<nb::tuple>(obj)) {
@@ -773,7 +777,7 @@ mlx_compute_scatter_args(
}
auto mlx_slice_update(
const array& src,
const mx::array& src,
const nb::object& obj,
const ScalarOrArray& v) {
// Can't route to slice update if not slice or tuple
@@ -784,7 +788,7 @@ auto mlx_slice_update(
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<array>(idx) || nb::isinstance<nb::list>(idx)) {
if (nb::isinstance<mx::array>(idx) || nb::isinstance<nb::list>(idx)) {
return std::make_pair(false, src);
}
}
@@ -881,7 +885,10 @@ auto mlx_slice_update(
return std::make_pair(true, out);
}
void mlx_set_item(array& src, const nb::object& obj, const ScalarOrArray& v) {
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);
@@ -897,8 +904,8 @@ void mlx_set_item(array& src, const nb::object& obj, const ScalarOrArray& v) {
}
}
array mlx_add_item(
const array& src,
mx::array mlx_add_item(
const mx::array& src,
const nb::object& obj,
const ScalarOrArray& v) {
auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
@@ -909,8 +916,8 @@ array mlx_add_item(
}
}
array mlx_subtract_item(
const array& src,
mx::array mlx_subtract_item(
const mx::array& src,
const nb::object& obj,
const ScalarOrArray& v) {
auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
@@ -921,8 +928,8 @@ array mlx_subtract_item(
}
}
array mlx_multiply_item(
const array& src,
mx::array mlx_multiply_item(
const mx::array& src,
const nb::object& obj,
const ScalarOrArray& v) {
auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
@@ -933,8 +940,8 @@ array mlx_multiply_item(
}
}
array mlx_divide_item(
const array& src,
mx::array mlx_divide_item(
const mx::array& src,
const nb::object& obj,
const ScalarOrArray& v) {
auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
@@ -945,8 +952,8 @@ array mlx_divide_item(
}
}
array mlx_maximum_item(
const array& src,
mx::array mlx_maximum_item(
const mx::array& src,
const nb::object& obj,
const ScalarOrArray& v) {
auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
@@ -957,8 +964,8 @@ array mlx_maximum_item(
}
}
array mlx_minimum_item(
const array& src,
mx::array mlx_minimum_item(
const mx::array& src,
const nb::object& obj,
const ScalarOrArray& v) {
auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);