Thanks @JuliaPoo! At first glance this looks quite good - the implementation is easy to understand, and docs, type annotation etc. all LGTM.

I experimented a little bit with it, and performance looks as expected: almost all the time is spent in np.argpartition and np.top_k is a lot faster than a full sort:

>>> x = np.random.rand(100_000)
>>> %timeit np.top_k(x, 5, largest=False)
759 µs ± 3.95 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
>>> %timeit np.argpartition(x, 5)
758 µs ± 3.86 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
>>> %timeit np.sort(x)[:5]
2.95 ms ± 3.29 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

So far so good. It doesn't seem like there'd be a significant performance or usability benefit, for this implementation at least, of adding more than one function (e.g., top_k_values).

I've drafted a PR to array-api-compat: data-apis/array-api-compat#158 that informs us about how compatible numpy's top_k implementation is to other existing implementations. Notably, it is directly compatible with torch.topk.

Are there any more concerns regarding this PR? Otherwise, I'll update the docstring to clarify the sort order of np.nan similar to #26716 and fix the merge conflict.

With #26716 merged, I've updated the docstring for top_k about the sort order of complex and nans, and added a release note.

Continuing from the discussion here, I've special-cased types in np.typecodes["AllFloat"] (float16, float32, float64, float64, complex64, complex128, complex128) to push np.nans to the back.

I.e., np.top_k(np.array([1,2,3,np.nan], dtype="f"), 3) will return values [1,2,3] instead of previously [2,3,np.nan].

Other array dtypes that could contain np.nans such as np.array([1,np.nan], dtype=object) will simply use the regular sort order (push nans to the front). The rational behind is that

• it's too difficult to support the nan behaviour for object arrays
• it's a case that is probably quite rare (ufuncs like np.isnan doesn't support object arrays anyways)
• for arbitruary dtypes such as that containing Decimals or a structured array, the onus should be on the user to handle np.nans

The sorting behavior is looking pretty good. I think the behavior for object arrays not being guaranteed is okay, that is very niche and can be documented as doing whatever np.sort does (same for custom dtypes). This does the right thing for integers and floating-point dtypes, which is a win.

Another question I'd like to raise is whether complex numbers should be supported at all. It doesn't seem useful, and we have a known issue with the semantics of sorting/comparing complex numbers being ill-defined. To illustrate:

>>> # 2+10j is the element with the largest absolute value
>>> x = np.asarray([1.5, 3, 4, 2+10j, np.nan, 9, 9], dtype=np.complex128)
>>> np.top_k(x, 3, largest=False)[0]
array([1.5 +0.j, 2. +10.j, 3.  +0.j])

For np.sort & co we've tended to not both trying to do anything about this, but I question whether we should be introducing support in a new function like this. If a user wants to use top_k on an array with a complex dtype, it seems better to have them use either top_k(x.real) or top_k(x.abs()) and use the returned indices to get the k complex values they want.

I personally found it a little weird to be sorting complex. Now that you mentioned it, I'm leaning towards np.top_k defaulting to np.sort ordering for complex.

We have discussed deprecating sorting complex before, and I still wouldn't mind it at all (ISTR an odd sort_complex function even for the odd-case where it is used).
Sort order is also used by minimum/maximum, so it seems fine to just do it and it would be fixed when the other places are... But I don't care if you want to explicitly disallow it either.