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 on the GPU, with no error raised and no crash.

Reproducer

import mlx.core as mx
import numpy as np

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

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

False before, True after.

Each output row is internally sorted but holds the wrong input row's data, so a check that only asks whether the output is sorted still passes. stream=mx.cpu is correct either way, as is any 2-D transpose. The failure needs a rank 3 or higher view whose non-sorted axes are no longer in row major order relative to each other.

Cause

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). Selection is on in.flags().contiguous (metal/sort.cpp:45, cuda/sort.cu:784), which only means dense storage with no gaps. That does not imply the non-sorted axes form a single ordered run, so a transposed view can take the contiguous kernel and sort the wrong rows.

Fix

Collapse the non-sorted axes with collapse_contiguous_dims and take the fast path only when they reduce to a single run, with the segment stride read from that run. The hand-rolled min-stride loop goes away.

This cannot demote an input that already works. The contiguous kernel is correct for a layout exactly when elem_to_loc(r) == r * stride holds for every r, and that is the arithmetic progression collapse_contiguous_dims merges, so every layout main gets right is still accepted.

Over 1058 (array, axis) cases, built from 15 base shapes of rank 1 to 4 by applying every axis permutation, expand_dims, broadcast_to, and forward, reversed and strided slices:

variant	fast + correct	fast + wrong	demoted
main	359	122	n/a
#4137 row_contiguous	301	0	58
this	359	0	0

The 122 layouts the selection predicate marks wrong on main are exactly the 122 whose GPU result differs from the CPU backend, with nothing in either set that is not in the other.

in.flags().contiguous stays in the condition. It is what keeps negative strides out of the contiguous kernel, whose row index is unsigned.

Testing

The new test covers (3, 4, 8), (2, 1, 6), (2, 3, 4, 2) and (2, 1, 3, 4) under every axis permutation and every axis. It fails on main and passes here; CPU is correct either way.

• Python: 898 passed, 5 skipped
• C++: 290/290 on GPU and on CPU
• pre-commit: clean

multi_block_sort is unaffected; multi block starts when the sorted axis exceeds 2048.

Benchmark

M3 Pro. Each entry is the median of 3 runs, each itself the median of 9 repeats x 50 iterations, in ms.

case	main	this
(1024,1024) axis 1, contiguous	0.4422	0.4496
(1024,1024) T(1,0) axis 0, transposed	0.4668	0.4725
(256,64,64) axis 2, contiguous	0.4856	0.4898
(64,64,256) T(1,0,2) axis 2	0.3636	0.4918
(512,512) axis 1, multi block	0.1950	0.1971

The first, second, third and fifth rows are unchanged within run to run spread; main alone varied from 0.4265 to 0.4587 on the first row across runs.

(64,64,256) T(1,0,2) is one of the layouts main sorts incorrectly, so this change moves it to the general path. That is the cost of correctness, not a regression on an input that worked. The transposed row above it is an input main sorts correctly, and it stays on the fast path.

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)