Collapse the non sorted axes to pick the contiguous sort kernel

[kapellirohith]

Supersedes #4137.

mx.sort and mx.argsort return silently wrong results on some transposed views: 242 of 9395 enumerated cases are wrong on main, with no error raised and no crash.

variant	fast+correct	fast+wrong	demoted
main	473	242	n/a
#4137 row_contiguous	374	0	99
this	473	0	0

2461 arrays, 9395 (array, axis) cases, GPU checked against the CPU backend, each variant built and run.

The enumeration builds those arrays from 15 base shapes of rank 1 to 4 by applying every axis permutation, expand_dims, broadcast_to, forward slices and reversed (negative stride) slices, then sorts each one on every axis. Per case it records which kernel the dispatch selects and compares the GPU result against the CPU backend, which does not use the segment stride. Selection is read from flags() and strides(), so it runs as a C++ doctest probe rather than a Python script; I can attach it if that is useful.

The contiguous kernel addresses row r at r * in_stride_segment_axis (kernels/sort.h:283, cuda/sort.cu:305). The nc kernel decodes the same row with elem_to_loc (kernels/sort.h:393, cuda/sort.cu:398). The kernel is selected on in.flags().contiguous (metal/sort.cpp:45, cuda/sort.cu:784), which means dense storage with no gaps and does not imply one ordered run, so a transposed view can take the contiguous kernel and sort the wrong rows.

Reproducer

Base 994d9d502
Chip M3 Pro
macOS 26.6.1

import mlx.core as mx
import numpy as np

a = np.random.RandomState(0).randn(3, 4, 8).astype(np.float32)
x = mx.array(a)
y = mx.swapaxes(x, 0, 1)

got = np.array(mx.sort(y, axis=-1))
ref = np.sort(np.swapaxes(a, 0, 1), axis=-1)

print(np.array_equal(got, ref))

Before: False
After: True

Each output row is internally sorted, but it holds the wrong input row's data, so a check that only asks whether the output is sorted will pass.

stream=mx.cpu is correct before and after. A 2-D transpose is also correct before and after. The failure is specific to the GPU single-block path on higher-rank transposed views where the non-sorted axes are no longer in row-major order relative to each other.

Fix

Collapse the non-sorted axes with collapse_contiguous_dims. The fast path is selected only when they reduce to a single run, with the segment stride read from that run, and the hand-rolled min-stride loop is removed. After collapsing, at most one non-size-1 dimension may remain, which is exactly the addressing assumption the contiguous kernel makes.

Testing

Confirmed the new tests fail on the main branch at 994d9d502: 5/5 red on GPU, both standalone and in-module. They pass on CPU either way; that backend never had the defect.

Shapes: (3, 4, 8), (2, 1, 6), (2, 3, 4, 2), (2, 1, 3, 4).

Python: 860 exit 0 on GPU. test_fft_too_large is a pre-existing CPU-only failure.
C++: 278/278 on both devices.
JIT: 10/10.
Pre-commit: clean.

multi_block_sort is unaffected. All 24 multi-block cases are correct, with multi-block starting when the sorted axis is greater than 2048.

Benchmark

case	main	this
(1024,1024) axis 1, contiguous	0.4515	0.4316
(1024,1024) T(1,0) axis 0, transposed	0.4719	0.4534
(256,64,64) axis 2, contiguous	0.4808	0.4691
(64,64,256) T(1,0,2) axis 2	0.3627	0.4675
(512,512) axis 1, multi block	0.1920	0.1867

Medians of 9 repeats x 50 iterations, 3 processes each.

The fast path is not broadly disabled: all 473 layouts that were fast and correct on main stay fast, and all 242 that were fast and wrong are rejected from it.

Build each revision with:

python setup.py build_ext --inplace

Run both revisions on the same GPU.

(64,64,256) T(1,0,2) is slower because main incorrectly selects the contiguous kernel for it; this change demotes it to the general path. That cost is the price of correctness, not a regression on a case that already worked. The transposed row above it is an input main sorts correctly, and it stays on the fast path.

Benchmark script:

import time
import statistics as st
import mlx.core as mx

mx.set_default_device(mx.gpu)


def bench(make, axis, reps=9, iters=50):
    med = []
    for _ in range(reps):
        a = make()
        mx.eval(a)
        mx.eval(mx.sort(a, axis=axis))
        t = time.perf_counter()
        for _ in range(iters):
            mx.eval(mx.sort(a, axis=axis))
        med.append((time.perf_counter() - t) / iters * 1e3)
    return st.median(med)


CASES = [
    ("(1024,1024) axis 1, contiguous", lambda: mx.random.normal((1024, 1024)), 1),
    ("(1024,1024) T(1,0) axis 0, transposed",
     lambda: mx.transpose(mx.random.normal((1024, 1024)), (1, 0)), 0),
    ("(256,64,64) axis 2, contiguous", lambda: mx.random.normal((256, 64, 64)), 2),
    ("(64,64,256) T(1,0,2) axis 2",
     lambda: mx.transpose(mx.random.normal((64, 64, 256)), (1, 0, 2)), 2),
    ("(512,512) axis 1, multi block", lambda: mx.random.normal((512, 512)), 1),
]

bench(CASES[0][1], CASES[0][2], reps=2)  # warm up, discarded
for name, make, axis in CASES:
    print(f"{name:42s} {bench(make, axis):8.4f} ms", flush=True)

CUDA validation is source-level only. No NVIDIA hardware was used.

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing this change
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)