perf(knn): replace argsort+slice with top_k for neighbor search, ~5-9x speedup

[Ljzn]

Problem

BruteKNN.brute_force_search/3 finds the K nearest neighbors by fully sorting
all distances and then slicing the first K:

Nx.argsort(distances, axis: 1, type: :u64) |> Nx.slice_along_axis(0, k, axis: 1)

This is O(N log N) per query — sorting all N distances just to retrieve K.
For N >> K, this is wasteful: the sort comparator runs ~N log N times while
only K elements are ultimately needed.

Solution

Replace with Nx.top_k/2 on negated distances, which selects the K smallest
values directly using XLA's native chlo.top_k operator (on EXLA) or a
block-based fallback:

{_values, raw_indices} = Nx.top_k(Nx.negate(distances), k: k)
neighbor_indices = Nx.as_type(raw_indices, {:u, 64})

This reduces the complexity from O(N log N) to O(N log K) per query.

Benchmark

BruteKNN.fit/predict on random 64-dimensional float data (EXLA backend),
averaged over 5 runs:

N (train)	Before (ms)	After (ms)	Speedup
5,000	0.77	0.15	5.1x
50,000	8.85	1.04	8.5x

Speedup increases with N/K ratio, consistent with the complexity class change.

Benchmark script

Mix.install([
  {:nx, path: "/path/to/nx", override: true},
  {:exla, path: "/path/to/exla"},
  {:scholar, path: "."}
], force: true, verbose: false)

Nx.global_default_backend(EXLA.Backend)
Nx.Defn.global_default_options(compiler: EXLA)

key = Nx.Random.key(42)
k = 5
for n <- [5000, 50000] do
  {d, key2} = Nx.Random.normal(key, 0.0, 1.0, shape: {n, 64}, type: {:f, 32})
  {q, _} = Nx.Random.normal(key2, 0.0, 1.0, shape: {1, 64}, type: {:f, 32})
  m = Scholar.Neighbors.BruteKNN.fit(d, num_neighbors: k)
  Scholar.Neighbors.BruteKNN.predict(m, q)
  times = Enum.map(1..5, fn _ ->
    {t, _} = :timer.tc(fn ->
      m2 = Scholar.Neighbors.BruteKNN.fit(d, num_neighbors: k)
      Scholar.Neighbors.BruteKNN.predict(m2, q)
    end); t
  end)
  IO.puts("n=#{n}\\t#{Float.round(Enum.sum(times) / length(times) / 1000, 2)} ms")
end

Correctness

Nx.top_k returns indices of the K largest values; negating the distances
converts "smallest K distances" into "largest K negative distances". The
resulting indices match those from argsort + slice.

Compatibility

This change is transparent:

• EXLA users benefit from the native chlo.top_k XLA operator
• BinaryBackend users fall back to argsort + take_along_axis + slice
  (same complexity as before, no regression for small K)
• All existing types (:u64 indices) are preserved via Nx.as_type

---

Assisted by DeepSeek V4 Flash

[krstopro]

Ah, good catch, well done!

[josevalim]

💚 💙 💜 💛 ❤️