Fix bugs with gather

Author: josevalim

exla/lib/exla/defn.ex

@@ -1285,15 +1285,17 @@ defmodule EXLA.Defn do
     tensor_shape = op_shape(tensor)
     tensor_rank = tuple_size(tensor_shape)
     tensor_axes = axes_for_rank(tensor_rank)
-    index_vector_dim = tuple_size(op_shape(indices)) - 1
+    indices_rank = tuple_size(op_shape(indices))
+    index_vector_dim = indices_rank - 1
 
     slice_sizes =
       for i <- tensor_axes do
         if i in axes, do: 1, else: elem(tensor_shape, i)
       end
 
     batch_size = tensor_rank - length(axes)
-    offset_dims = count_up(batch_size, batch_size)
+    offset_size = indices_rank - length(axes)
+    offset_dims = count_up(batch_size, offset_size)
 
     Value.gather(
       tensor,

nx/lib/nx.ex

@@ -7876,6 +7876,8 @@ defmodule Nx do
   end
 
   defp indexed_axes(tensor, indices, opts) do
+    n = elem(indices.shape, tuple_size(indices.shape) - 1)
+
     if axes = opts[:axes] do
       axes = Nx.Shape.normalize_axes(tensor.shape, axes, tensor.names)
 
@@ -7884,9 +7886,13 @@ defmodule Nx do
         _, _ -> raise ArgumentError, ":axes must be an ordered list"
       end)
 
+      if length(axes) != n do
+        raise ArgumentError,
+              ":axes must have the same number of elements as the last dimension of indices"
+      end
+
       axes
     else
-      n = elem(indices.shape, tuple_size(indices.shape) - 1)
       Enum.to_list(0..(n - 1))
     end
   end