Sort once

Every problem in this track so far has been about making an operation cheaper. This one is about not doing it.

def rank_of(values, queries):
    out = []
    for q in queries:
        s, _ = values.sort()            # once per query
        out.append((s < q).sum())
    return torch.stack(out)

The sort does not depend on q. It is the hoisting problem again, and it is worth its own place because the thing being hoisted is asymptotically expensive rather than merely allocating.

Sorting n items is n log n. Doing it once per query makes the whole routine q * n log n where it should be n log n + q log n, and no amount of kernel-level tuning recovers that.

Then delete the loop as well

Hoisting leaves a Python loop over queries. searchsorted does all of them at once against the sorted array:

s, _ = values.sort()
return torch.searchsorted(s, queries)

Two operations, whatever the number of queries. This is the loop-deletion lesson and the hoisting lesson composing, which is what most real optimisation looks like: not one trick, but the same two or three applied until nothing is left.

The order to work in

1. Algorithm. Is this the right amount of work? Sorting once beats sorting q times by more than any kernel choice.
2. Vectorisation. Is it one kernel or n?
3. Memory. Is it allocating more than it needs?
4. Fusion and compilation. Only now.

Reaching for step 4 while step 1 is wrong is the most common way to spend a day and gain nothing.

Your task

def ranks(values: torch.Tensor, queries: torch.Tensor) -> torch.Tensor

For each query, return how many values are strictly less than it, sorting at most once and using at most four dispatched operations regardless of how many queries there are.