Nx.take: support indices of rank > 1

[grzuy]

candlex/test/candlex_test.exs

Lines 2203 to 2218 in 9e133fe

	# t([[1, 2], [11, 12]])
	# |> Nx.take(t([[0, 0], [1, 1], [0, 0]]), axis: 1)
	# |> assert_equal(t(
	# [
	# [
	# [1, 1],
	# [2, 2],
	# [1, 1]
	# ],
	# [
	# [11, 11],
	# [12, 12],
	# [11, 11]
	# ]
	# ]
	# ))

[grzuy]

Consider elixir-nx/nx#1366.