mlx_median with no axis (axes=NULL, axes_num=0) fails: [flatten] start_axis must be less than or equal to end_axis

[aki-xavier]

Describe the bug

mlx_median has no dedicated no-axis entry point, so the "median of all elements" case is expressed as axes = NULL, axes_num = 0. That path builds std::vector<int>(axes, axes + 0) — an empty axes vector — which flows into the axes-overload of mlx::core::median.

With no median axes, transpose_axes = [0..ndim) and flat_start = ndim, so it calls flatten(..., start_axis = ndim, end_axis = ndim). The flatten clamp end_ax = min(ndim-1, ndim) = ndim-1 < start_ax triggers the assertion [flatten] start_axis must be less than or equal to end_axis.

Every other reduction (mlx_sum / mlx_mean / mlx_max / mlx_min / mlx_var / mlx_logsumexp) has a dedicated no-axis C entry point that maps directly to the core full-reduce overload. mlx_median is the sole exception, which is why only median is broken.

Repro (C API)

mlx_array res = NULL;
mlx_array a = /* e.g. a 2x2 float32 array */;
int rc = mlx_median(&res, a, NULL /* axes */, 0 /* axes_num */, false, NULL /* default stream */);
// rc == 1
// error = "[flatten] start_axis must be less than or equal to end_axis at mlx/c/ops.cpp:2115"

mlx_median(&res, a, NULL, 0, false, s) is the documented way to reduce over all axes.

Expected vs. actual

• Expected: rc == 0 and res holds the median of all elements.
• Actual: rc == 1 with the [flatten] error above, so the no-axis median is unusable.

Proposed fix

Special-case axes_num == 0 inside mlx_median and route it to the full-reduce overload mlx::core::median(a, keepdims, s), mirroring the other reductions:

extern "C" int mlx_median(
    mlx_array* res,
    const mlx_array a,
    const int* axes,
    size_t axes_num,
    bool keepdims,
    const mlx_stream s) {
  try {
    // axes_num == 0 means "no axis": reduce over every dimension. The axes
    // overload of mlx::core::median does not accept an empty axes list (it
    // feeds flatten(start == ndim), which asserts), so route the no-axis case
    // to the dedicated full-reduce overload, like mlx_sum / mlx_mean etc.
    if (axes_num == 0) {
      mlx_array_set_(
          *res,
          mlx::core::median(mlx_array_get_(a), keepdims, mlx_stream_get_(s)));
    } else {
      mlx_array_set_(
          *res,
          mlx::core::median(
              mlx_array_get_(a),
              std::vector<int>(axes, axes + axes_num),
              keepdims,
              mlx_stream_get_(s)));
    }
  } catch (std::exception& e) {
    mlx_error(e.what());
    return 1;
  }
  return 0;
}

The full-reduce overload mlx::core::median(const array&, bool keepdims, StreamOrDevice) already exists in mlx/ops.h, so this compiles unchanged.

Environment

• mlx-c 0.6.0 (0.6.0_4)
• core mlx 0.32.1
• macOS Apple Silicon