Use the row contiguous flag to pick the contiguous sort kernel - #4137

Problem

Sorting a transposed view returns wrong results on the GPU. The values inside
each row are sorted correctly, only the rows land in the wrong places, so
nothing raises and the output looks plausible.

Minimal reproducer (M3 Pro, macOS 26.x, main @ e78d894):

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((got != ref).sum(), "/", got.size)   # before: 80 / 96 ; after: 0 / 96

stream=mx.cpu: 0 / 96. A 2-D transpose: 0 / 96. mx.transpose(x, (1, 2, 0)):
0 / 96. Only the GPU, only 3-D and higher, and only when the axes that are not
sorted end up out of row major order relative to each other.

Mechanism

single_block_sort selects the contiguous kernel on in.flags().contiguous,
which only means the data is dense, not that it is row major. A transpose keeps
that flag. The contiguous kernel enumerates rows with a single segment stride,
the smallest stride among the axes that are not sorted, and
Sort::eval_gpu allocates the output with a plain
allocator::malloc(out.nbytes()), so the output is always row major while the
input is permuted. The kernel then reads the rows in one order and writes them
in another.

Positive identification: with a of shape (3, 4, 8) and y = swapaxes(a, 0, 1),
y has strides (8, 32, 1). Removing the sorted axis leaves strides (8, 32)
over shape (4, 3), whose minimum is 8, so row r is read at offset 8r. The
correct offset for logical row (i, j) is 8i + 32j. Every output row of the
observed run is a correctly sorted row of the input, just the wrong one:
np.sort(got, axis=None) equals np.sort(ref, axis=None) exactly, and each
output row matches the input row at 8r rather than at 8i + 32j.

check_strides already restricts the fast path to a sorted axis holding the
smallest or largest stride. For a row major input those are the only two cases,
and both enumerate correctly, which is why this needs a permuted view to fire.
multi_block_sort passes the full strides and is unaffected, so the sorted axis
also has to be short enough for the single block path.

Introduced by 2615660 "Fix strided sort bug (#1236)" (2024-06-26), which added
output strides to the kernel. This does not revert that, it narrows when the
fast path is selected.

The CUDA backend has the identical line, the identical check_strides, the
identical out.set_data(cu::malloc_async(out.nbytes(), encoder)), and the same
inp += blockIdx.y * in_stride_segment_axis row enumeration, so it is changed
too. I have no NVIDIA GPU. That half is reasoned from source and validated by
CI, not run locally.

Why existing tests missed it

test_sort covers strided inputs via a_mx[::2, :, ::2], which keeps the axes
in row major order, and every other sort test uses freshly created arrays. No
sort test transposed its input.

Perf

row_contiguous is sufficient but not necessary. A permutation such as
(1, 2, 0) puts the sorted axis on the largest stride and leaves the other axes
in row major order relative to each other, so it was correct before and now
takes the general kernel:

Medians of 9 repeats of 100 iterations, population stdev 0.001 to 0.013 ms. A
row major sort of the same shape, measured in the same runs as a control, moves
by at most 5 percent between the two builds, which bounds the drift between
them. Row major inputs take the same path as before.

I posted "the measured cost is zero" in an earlier revision of this description.
That was wrong. Those numbers were taken on perm(1, 0, 2), which was already
returning wrong results before this change and so was never a deoptimisation.
The perm(1, 2, 0) family above is the set that was correct before and is
slower now, and the cost there is 5 to 11 percent.

A tighter predicate that keeps the fast path whenever the axes that are not
sorted are packed among themselves recovers all of it (0.863, 0.522, 1.031,
0.154 ms) at about 18 lines against one word. I have it written and tested and
will send it if you prefer, but I would rather land the small correct version
first.

Validation

M3 Pro, macOS 26.x, against main @ e78d894.

The 510 differing arrays were each classified against a NumPy reference: all
510 were wrong on main and correct here, none went the other way, and none
changed while both were correct, so no tie-breaking behavior moved.

Not verified locally: CUDA, Linux, Windows. The cuda/sort.cu change is
source-level reasoning validated by CI.

Checklist

Put an x in the boxes that apply.

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