Use stable sorting on KDTrees and fix warnings on MDS

msluszniak · Dec 11, 2023 · acb0673 · acb0673
1 parent 6845727
commit acb0673
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/lib/scholar/manifold/mds.ex b/lib/scholar/manifold/mds.ex
@@ -328,7 +328,6 @@ defmodule Scholar.Manifold.MDS do
   ## Examples
 
       iex> x = Nx.iota({4,5})
-      iex> key = Nx.Random.key(42)
       iex> init = Nx.reverse(Nx.iota({4,2}))
       iex> Scholar.Manifold.MDS.fit(x, init)
       %Scholar.Manifold.MDS{

diff --git a/lib/scholar/neighbors/kd_tree.ex b/lib/scholar/neighbors/kd_tree.ex
@@ -188,13 +188,13 @@ defmodule Scholar.Neighbors.KDTree do
     {level, tags, _tensor, _band} =
       while {level = Nx.u32(0), tags, tensor, band}, level < levels - 1 do
         k = rem(level, dims)
-        indices = Nx.argsort(tensor[[.., k]] + band * tags, type: :u32)
+        indices = Nx.argsort(tensor[[.., k]] + band * tags, type: :u32, stable: true)
         tags = update_tags(tags, indices, level, levels, size)
         {level + 1, tags, tensor, band}
       end
 
     k = rem(level, dims)
-    Nx.argsort(tensor[[.., k]] + band * tags, type: :u32)
+    Nx.argsort(tensor[[.., k]] + band * tags, type: :u32, stable: true)
   end
 
   defnp update_tags(tags, indices, level, levels, size) do