feat: gather

c_src/emlx_nif.cpp

@@ -437,7 +437,8 @@ NIF(pad) {
   TENSOR_PARAM(4, pad_value);
   DEVICE_PARAM(5, device);
 
-  TENSOR(mlx::core::pad(*t, axes, low_pad_size, high_pad_size, *pad_value, "constant", device))
+  TENSOR(mlx::core::pad(*t, axes, low_pad_size, high_pad_size, *pad_value,
+                        "constant", device))
 };
 
 NIF(sort) {
@@ -505,6 +506,16 @@ NIF(take) {
   TENSOR(mlx::core::take(*t, *indices, axis, device));
 }
 
+NIF(gather) {
+  TENSOR_PARAM(0, t);
+  LIST_PARAM(1, std::vector<mlx::core::array>, indices);
+  LIST_PARAM(2, std::vector<int>, axes);
+  LIST_PARAM(3, std::vector<int>, slice_sizes);
+  DEVICE_PARAM(4, device);
+
+  TENSOR(mlx::core::gather(*t, indices, axes, slice_sizes, device));
+}
+
 /* Reduction Ops */
 
 #define REDUCTION_AXES_OP(OP) REDUCTION_AXES_OP2(OP, OP)
@@ -877,6 +888,7 @@ static ErlNifFunc nif_funcs[] = {{"strides", 1, strides},
                                  {"concatenate", 3, concatenate},
                                  {"take_along_axis", 4, take_along_axis},
                                  {"take", 4, take},
+                                 {"gather", 5, gather},
                                  {"slice", 5, slice},
                                  {"slice_update", 5, slice_update},
                                  {"squeeze", 3, squeeze},

lib/emlx.ex

@@ -231,6 +231,7 @@ defmodule EMLX do
   deftensor concatenate(tensors, axis)
   deftensor take_along_axis(tensor, tensorIndices, axis)
   deftensor take(tensor, tensorIndices, axis)
+  deftensor gather(tensor, indices, axes, slice_sizes)
   deftensor max(tensor, axes, keep_axes)
   deftensor min(tensor, axes, keep_axes)
   deftensor clip(tensor, tensor_min, tensor_max)

lib/emlx/backend.ex

@@ -1114,6 +1114,39 @@ defmodule EMLX.Backend do
     end)
   end
 
+  @impl true
+  def gather(out, tensor, indices, opts) do
+    axes = opts[:axes]
+
+    num_axes = Nx.axis_size(indices, -1)
+
+    slice_sizes =
+      Enum.map(Nx.axes(tensor), fn axis ->
+        if axis in axes do
+          1
+        else
+          Nx.axis_size(tensor, axis)
+        end
+      end)
+
+    indices_list =
+      Enum.map(0..(num_axes - 1), fn entry ->
+        {_device, ref} =
+          indices
+          |> Nx.slice_along_axis(entry, 1, axis: -1)
+          |> Nx.squeeze(axes: [-1])
+          |> from_nx()
+
+        ref
+      end)
+
+    tensor
+    |> from_nx()
+    |> EMLX.gather(indices_list, axes, slice_sizes)
+    |> EMLX.reshape(out.shape)
+    |> to_nx(out)
+  end
+
   for {op, arity} <- [
         reduce: 5,
         window_reduce: 6,
@@ -1135,7 +1168,6 @@ defmodule EMLX.Backend do
         to_pointer: 2,
         indexed_put: 5,
         indexed_add: 5,
-        gather: 4,
         from_pointer: 5
       ] do
     @impl true

lib/emlx/nif.ex

@@ -189,6 +189,10 @@ defmodule EMLX.NIF do
     :erlang.nif_error(:nif_not_loaded)
   end
 
+  def gather(_tensor, _indices, _axes, _slice_sizes, _device) do
+    :erlang.nif_error(:nif_not_loaded)
+  end
+
   def abs(_tensor, _device) do
     :erlang.nif_error(:nif_not_loaded)
   end

test/emlx/nx_doctest_test.exs

@@ -11,8 +11,6 @@ defmodule EMLX.Nx.DoctestTest do
     indexed_put: 4,
     # put_diagonal depends on indexed_put
     put_diagonal: 3,
-    # take_diagonal depends on gather
-    take_diagonal: 2,
     # make_diagonal depends on indexed_put
     make_diagonal: 2,
     # mode depends on indexed_add
@@ -26,9 +24,7 @@ defmodule EMLX.Nx.DoctestTest do
     window_scatter_min: 5,
     window_scatter_max: 5,
     reverse: 2,
-    take: 3,
     slice: 4,
-    gather: 3,
     reflect: 2
   ]