Skip to content

Replace hand-rolled axis normalization with normalize_axis_index in split/unstack/partition/topk #4288

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zcbenz merged 1 commit into from
Aug 16, 2026
Merged

Replace hand-rolled axis normalization with normalize_axis_index in split/unstack/partition/topk #4288

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 8 additions & 46 deletions mlx/ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1143,13 +1143,7 @@ std::vector<array> split(
const Shape& indices,
int axis,
StreamOrDevice s /* = {} */) {
auto ax = axis < 0 ? axis + a.ndim() : axis;
if (ax < 0 || ax >= a.ndim()) {
std::ostringstream msg;
msg << "Invalid axis (" << axis << ") passed to split"
<< " for array with shape " << a.shape() << ".";
throw std::invalid_argument(msg.str());
}
auto ax = normalize_axis_index(axis, a.ndim(), "[split] ");

if (indices.empty()) {
return {a};
Expand Down Expand Up @@ -1191,20 +1185,14 @@ split(const array& a, const Shape& indices, StreamOrDevice s /* = {} */) {

std::vector<array>
split(const array& a, int num_splits, int axis, StreamOrDevice s /* = {} */) {
auto ax = axis < 0 ? axis + a.ndim() : axis;
if (ax < 0 || ax >= a.ndim()) {
std::ostringstream msg;
msg << "Invalid axis " << axis << " passed to split"
<< " for array with shape " << a.shape() << ".";
throw std::invalid_argument(msg.str());
}
auto ax = normalize_axis_index(axis, a.ndim(), "[split] ");
if (num_splits <= 0) {
std::ostringstream msg;
msg << "[split] num_splits must be positive and non-zero but got "
<< num_splits << ".";
throw std::invalid_argument(msg.str());
}
auto q_and_r = std::ldiv(a.shape(axis), num_splits);
auto q_and_r = std::ldiv(a.shape(ax), num_splits);
if (q_and_r.rem) {
std::ostringstream msg;
msg << "Array split does not result in sub arrays with equal size:"
Expand All @@ -1217,7 +1205,7 @@ split(const array& a, int num_splits, int axis, StreamOrDevice s /* = {} */) {
for (int i = 0; i < indices.size(); ++i) {
indices[i] = (i + 1) * split_size;
}
return split(a, indices, axis, s);
return split(a, indices, ax, s);
}

std::vector<array>
Expand All @@ -1228,13 +1216,7 @@ split(const array& a, int num_splits, StreamOrDevice s /* = {} */) {
std::vector<array>
unstack(const array& a, int axis, StreamOrDevice s /* = {} */) {
auto ndim = static_cast<int>(a.ndim());
auto ax = axis < 0 ? axis + ndim : axis;
if (ax < 0 || ax >= ndim) {
std::ostringstream msg;
msg << "[unstack] Invalid axis " << axis << " for array with " << ndim
<< " dimensions.";
throw std::invalid_argument(msg.str());
}
auto ax = normalize_axis_index(axis, ndim, "[unstack] ");
auto n = a.shape(ax);
std::vector<array> res;
res.reserve(n);
Expand Down Expand Up @@ -2838,14 +2820,7 @@ array partition(
int axis,
StreamOrDevice s /* = {} */) {
// Check for valid axis
if (axis + static_cast<int>(a.ndim()) < 0 ||
axis >= static_cast<int>(a.ndim())) {
std::ostringstream msg;
msg << "[partition] Received invalid axis " << axis << " for array with "
<< a.ndim() << " dimensions.";
throw std::invalid_argument(msg.str());
}
int axis_ = axis < 0 ? axis + a.ndim() : axis;
int axis_ = normalize_axis_index(axis, a.ndim(), "[partition] ");
int kth_ = kth < 0 ? kth + a.shape(axis) : kth;
if (kth_ < 0 || kth_ >= a.shape(axis_)) {
std::ostringstream msg;
Expand Down Expand Up @@ -2879,14 +2854,7 @@ array argpartition(
int axis,
StreamOrDevice s /* = {} */) {
// Check for valid axis
if (axis + static_cast<int>(a.ndim()) < 0 ||
axis >= static_cast<int>(a.ndim())) {
std::ostringstream msg;
msg << "[argpartition] Received invalid axis " << axis << " for array with "
<< a.ndim() << " dimensions.";
throw std::invalid_argument(msg.str());
}
int axis_ = axis < 0 ? axis + a.ndim() : axis;
int axis_ = normalize_axis_index(axis, a.ndim(), "[argpartition] ");
int kth_ = kth < 0 ? kth + a.shape(axis) : kth;
if (kth_ < 0 || kth_ >= a.shape(axis_)) {
std::ostringstream msg;
Expand Down Expand Up @@ -2941,13 +2909,7 @@ array topk(const array& a, int k, StreamOrDevice s /* = {}*/) {
/** Returns topk elements of the array along a given axis. */
array topk(const array& a, int k, int axis, StreamOrDevice s /* = {}*/) {
// Check for valid axis
int axis_ = axis < 0 ? axis + a.ndim() : axis;
if (axis_ < 0 || axis_ >= static_cast<int>(a.ndim())) {
std::ostringstream msg;
msg << "[topk] Received invalid axis " << axis << " for array with "
<< a.ndim() << " dimensions.";
throw std::invalid_argument(msg.str());
}
int axis_ = normalize_axis_index(axis, a.ndim(), "[topk] ");
if (k < 0 || k > a.shape(axis_)) {
std::ostringstream msg;
msg << "[topk] Received invalid k=" << k << " along axis " << axis
Expand Down
Loading