Cast multiple indexes on slice, closes #1472

Author: josevalim

exla/lib/exla/defn.ex

@@ -1184,6 +1184,15 @@ defmodule EXLA.Defn do
       limit_indices = Enum.zip_with(start_indices, lengths, fn i, len -> i + len end)
       Value.slice(tensor, start_indices, limit_indices, strides)
     else
+      sample = Enum.find(start_indices, &(not is_integer(&1)))
+
+      type =
+        Enum.reduce(start_indices, op_type(sample), fn
+          index, acc when is_integer(index) -> acc
+          value, acc -> merge_type(op_type(value), acc)
+        end)
+
+      start_indices = Enum.map(start_indices, &to_type(&1, type))
       zeros = List.duplicate(0, tuple_size(ans.shape))
       slice = Value.dynamic_slice(tensor, start_indices, lengths)
 

exla/test/exla/backend_test.exs

@@ -181,6 +181,16 @@ defmodule EXLA.BackendTest do
     end
   end
 
+  describe "access" do
+    test "multiple indexes" do
+      tensor = Nx.eye({4, 4})
+      index = Nx.u32(2)
+      swap = Nx.s64(0)
+      assert tensor[[index, swap]] |> Nx.to_number() == 0
+      assert tensor[[0, swap]] |> Nx.to_number() == 1
+    end
+  end
+
   test "conjugate" do
     assert inspect(Nx.conjugate(~V[1 2-0i 3+0i 0-i 0-2i])) =~
              "1.0-0.0i, 2.0+0.0i, 3.0-0.0i, 0.0+1.0i, 0.0+2.0i"