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

[Adityaj0]

Fixes #4287.

Proposed changes

split, unstack, partition, argpartition and topk each duplicated the same axis normalization pattern:

auto ax = axis < 0 ? axis + a.ndim() : axis;
if (ax < 0 || ax >= a.ndim()) {
  std::ostringstream msg;
  msg << "...";
  throw std::invalid_argument(msg.str());
}

with slightly different bounds-check expressions at each site. normalize_axis_index already does this — normalize and validate in one call — and is already used for exactly this purpose elsewhere in ops.cpp (squeeze, expand_dims, flip, concatenate, stack, repeat). #4118 fixed the same duplication in take_along_axis/scatter_axis/linalg::cross; this covers the remaining call sites in ops.cpp.

This is a pure refactor, not a behavior change:

• Valid axes, including negative ones, normalize to the same value as before.
• Out-of-range axes still raise std::invalid_argument (ValueError in Python).
• Only the exact wording of the error message changes (e.g. [topk] Axis -100 is out of bounds for array with 3 dimensions. instead of [topk] Received invalid axis -100 for array with 3 dimensions.); no test in the tree asserts on that text, only that the correct exception type is raised.

In split(a, num_splits, axis, s) the normalized axis is now also reused for the subsequent a.shape(...) lookup and the recursive split(a, indices, ax, s) call instead of re-deriving it, avoiding the redundant re-normalization that happened implicitly through array::shape(int).

Net diff removes more than it adds (54 lines changed, 8 added / 46 removed).

Tests

No behavior change, so no new tests were added. Verified with a build of libmlx.a (cmake --build cmake_build --target mlx) and the full C++ test suite (cmake --build cmake_build --target tests && ./cmake_build/tests/tests): 263 test cases / 3578 assertions, all passing, including ops_tests.cpp coverage of split, unstack, partition, argpartition and topk.

• I have read the CONTRIBUTING document
• I have added tests that prove my fix is effective or that my feature works (n/a — pure refactor, existing tests cover behavior)
• I have updated the necessary documentation (not needed, no API or behavior change)
• I have run pre-commit run --all-files to format my code and installed pre-commit prior to committing changes