test: Nx.slice/3 with tensors as start_indices

mimiquate · Nov 17, 2023 · 9677568 · 9677568
1 parent c13ab7e
commit 9677568
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/test/candlex_test.exs b/test/candlex_test.exs
@@ -1824,6 +1824,18 @@ defmodule CandlexTest do
       #     [5, 6]
       #   ]
       # ))
+
+      t([0, 1])
+      |> Nx.slice([t(0)], [1])
+      |> assert_equal(t([0]))
+
+      t([0, 1])
+      |> Nx.slice([t(1)], [1])
+      |> assert_equal(t([1]))
+
+      t([[1, 2, 3], [4, 5, 6]])
+      |> Nx.slice([t(0), t(1)], [1, 1])
+      |> assert_equal(t([[2]]))
     end
 
     test "squeeze" do