Add mx.searchsorted

[erwinzhang7]

Closes #1255.

Adds mx.searchsorted(sorted_sequence, values, side="left"), matching np.searchsorted semantics.

Approach

Composed from existing primitives in ops.cpp rather than added as a new primitive. That means no new Metal or CUDA kernels, and it works on every backend immediately. It is also the approach originally suggested in the issue.

This is deliberately different from the earlier attempt in #2817, which added a primitive with a CPU implementation, left GPU evaluation as a TODO, and came to 901 lines across 13 files before being closed. The whole change here is 170 lines including tests and docs.

Which algorithm, and the dispatch

The issue left an open question: whether to use the linear form, the binary search, or dispatch between them. The suggestion at the time was to measure and pick. Here are the numbers on an M5 Max, current main, bf16 aside and all float32:

sorted n	queries m	binary	binary compiled	linear
1024	1	0.323	0.259	0.159
1024	1024	0.335	0.259	0.169
16384	16384	0.389	0.283	1.434
262144	16	0.436	0.290	0.177
262144	1024	0.441	0.314	1.453
2097152	1	0.440	0.319	0.259
2097152	1024	0.470	0.332	15.039
2097152	16384	0.513	0.346	out of memory

Milliseconds. The binary search is close to flat across a 2000x range in n, while the linear form is cheaper at small sizes and then degrades badly and eventually cannot allocate. The crossover sits between 4M and 33M for the n * m product, so the implementation dispatches at 4M.

That gives the better of the two everywhere. Against the workaround people are currently copying out of the issue thread:

sorted n	queries m	branch taken	mx.searchsorted	workaround	speedup
1024	64	linear	0.166	0.162	0.98x
16384	256	linear	0.177	0.177	1.00x
262144	16	linear	0.177	0.176	1.00x
262144	1024	binary	0.440	1.465	3.33x
2097152	256	binary	0.437	2.781	6.37x
2097152	1024	binary	0.457	15.328	33.5x
2097152	16384	binary	0.472	240.096	509x

Below the threshold it is the same computation, so parity is expected. Note the last row allocates a 34 billion element mask for the workaround, which only completes at all because this machine has 128 GB.

All timings warm up and measure by wall clock rather than by iteration count. With a fixed iteration count these kernels are short enough that the measurement tracks the GPU clock ramp instead of the kernel, and the same config can drift by 3x between processes.

Implementation note

The binary search accumulates into a single running result:

res = 0
for step in [2^(k-1) ... 2, 1]:
    cand = res + step
    ok   = (cand <= n) and (a[cand-1] < v)
    res  = where(ok, cand, res)

rather than tracking a lower and upper bound. The pair form needs an extra guard, because once lower and upper converge the midpoint update can push lower past upper, and the gather can index n. Accumulating avoids both without the extra mask.

Semantics

Validated against np.searchsorted on 178 cases before writing any C++, then again through the built op on 44 more, forcing both dispatch branches and running on both the GPU and CPU streams. Covers both sides, duplicates, exact matches, values below and above the range, n = 1, integer and mixed dtypes, and multi dimensional values, plus an exhaustive check for every n from 1 to 32 against every gap and every exact value.

Returns uint32 to match argsort, argmax and argpartition.

Scope

sorted_sequence must be 1-D, which matches np.searchsorted. Batched sorted sequences, as torch.searchsorted supports, are not included here. That would slot in as a take_along_axis variant if wanted, but I kept this to the shape the issue asked for rather than widening it.

[zcbenz]

Thanks for the PR but this is something that we want to have a kernel rather than a fallback. Sometimes we do add fallback implementations when it is absolutely necessary, but I don't think this is the case.

[erwinzhang7]

Noted, thanks for taking a look.

Working on a real kernel version instead, with CPU, Metal and CUDA implementations. Will open a fresh PR once it is ready rather than reviving this one.