move get_batches inside BruteKNN

lib/scholar/neighbors/brute_knn.ex

@@ -212,6 +212,27 @@ defmodule Scholar.Neighbors.BruteKNN do
     {neighbor_indices, neighbor_distances}
   end
 
+  defn get_batches(tensor, opts) do
+    {size, dim} = Nx.shape(tensor)
+    batch_size = opts[:batch_size]
+    num_batches = div(size, batch_size)
+    leftover_size = rem(size, batch_size)
+
+    batches =
+      tensor
+      |> Nx.slice_along_axis(0, num_batches * batch_size, axis: 0)
+      |> Nx.reshape({num_batches, batch_size, dim})
+
+    leftover =
+      if leftover_size > 0 do
+        Nx.slice_along_axis(tensor, num_batches * batch_size, leftover_size, axis: 0)
+      else
+        nil
+      end
+
+    {batches, leftover}
+  end
+
   defnp brute_force_search(data, query, opts) do
     k = opts[:num_neighbors]
     metric = opts[:metric]

lib/scholar/shared.ex

@@ -89,25 +89,4 @@ defmodule Scholar.Shared do
 
     valid_broadcast(to_parse - 1, n_dims, shape1, shape2)
   end
-
-  defn get_batches(tensor, opts) do
-    {size, dim} = Nx.shape(tensor)
-    batch_size = opts[:batch_size]
-    num_batches = div(size, batch_size)
-    leftover_size = rem(size, batch_size)
-
-    batches =
-      tensor
-      |> Nx.slice_along_axis(0, num_batches * batch_size, axis: 0)
-      |> Nx.reshape({num_batches, batch_size, dim})
-
-    leftover =
-      if leftover_size > 0 do
-        Nx.slice_along_axis(tensor, num_batches * batch_size, leftover_size, axis: 0)
-      else
-        nil
-      end
-
-    {batches, leftover}
-  end
 end