diff --git a/mlx/ops.cpp b/mlx/ops.cpp
index d154bd3d19..a0354fabf8 100644
--- a/mlx/ops.cpp
+++ b/mlx/ops.cpp
@@ -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};
@@ -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:"
@@ -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>
@@ -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);
@@ -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;
@@ -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;
@@ -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