Replace second order argsort with permutation inverse (elixir-nx#223)

lib/scholar/neighbors/kd_tree.ex

@@ -82,7 +82,7 @@ defmodule Scholar.Neighbors.KDTree do
   end
 
   defnp update_tags(tags, indices, level, levels, size) do
-    pos = Nx.argsort(indices, type: :u32)
+    pos = inverse_permutation(indices)
 
     pivot =
       bounded_segment_begin(tags, levels, size) +
@@ -103,6 +103,17 @@ defmodule Scholar.Neighbors.KDTree do
     )
   end
 
+  defnp inverse_permutation(indices) do
+    shape = Nx.shape(indices)
+    type = Nx.type(indices)
+
+    Nx.indexed_put(
+      Nx.broadcast(Nx.tensor(0, type: type), shape),
+      Nx.new_axis(indices, -1),
+      Nx.iota(shape, type: type)
+    )
+  end
+
   defnp bounded_subtree_size(i, levels, size) do
     diff = levels - bounded_level(i) - 1
     shifted = 1 <<< diff