Often times you want to slice and index into specific parts of a tensor. Nx offers a few different slicing and indexing routines which allow you to accomplish most of what you would want to do. Slicing can be a bit tricky given static shape requirements, but you usually can work around limitations.

First, you can perform generic slices using Nx.slice/4:

a = Nx.tensor([[1, 2, 3], [4, 5, 6]])
Nx.slice(a, [0, 1], [1, 2])

Returns:

#Nx.Tensor<
  s64[1][2]
  [
    [2, 3]
  ]
>

The first list in Nx.slice/4 is a list of start indices. You may specify dynamic or tensor values for the start index - as long as they have a scalar shape:

Nx.slice(a, [Nx.tensor(0), Nx.tensor(1)], [1, 2])

The second list is the length of each slice. Each value must be known or static at compile-time. This is because the length dictates the final shape of the sliced tensor:

Nx.slice(a, [1, 2], [2, 3])

Returns:

#Nx.Tensor<
  s64[2][3]
  [
    [1, 2, 3],
    [4, 5, 6]
  ]
>

This might surprise you! You might have expected an out-of-bounds error or something similar, considering that you are requesting slices from both axes past the original bounds of the tensor. In reality, Nx will force the slice in bounds to match the length requested from the slice in each dimension.

You can also make use of the Access syntax which builds on top of normal slice operations. For example:

a = Nx.tensor([[1, 2, 3], [4, 5, 6]])
a[[1, 0..1]]

Returns:

#Nx.Tensor<
  s64[2]
  [4, 5]
>

Because you’re accessing the first index from the first dimension and the zero to first indices from the second dimension. In newer versions of Elixir, you can slice an entire axis with ..:

a[[.., 2]]

#Nx.Tensor<
  s64[2]
  [3, 6]
>

If Nx.slice/4 and the Access syntax is not flexible enough for you, you can try one of Nx.take/3, Nx.gather/3, and Nx.take_along_axis/3. For example, a common operation in deep learning is computing a vector representation from sparse values. Given a sequence of integer tokens between 0 and 127, you can convert each token to a vector representation using Nx.take/3:

tokens = Nx.tensor([127, 32, 0, 1, 5, 6])
weights = Nx.random_uniform({128, 32})
Nx.take(weights, tokens)

Returns:

#Nx.Tensor<
  f32[6][32]
  [
    [0.39863643050193787, 0.33112287521362305, 0.531199038028717, 0.3594178259372711, 0.8754940629005432, 0.30342867970466614, 0.9188190698623657, 0.29185304045677185, 0.543312668800354, 0.5064964294433594, 0.7225326299667358, 0.06837604194879532, 0.5449554920196533, 0.2207975834608078, 0.0635833740234375, 0.3370073139667511, 0.6428131461143494, 0.8821378946304321, 0.9932462573051453, 0.8975431323051453, 0.7079696655273438, 0.023084526881575584, 0.4048435091972351, 0.12792034447193146, 0.4222281277179718, 0.21171192824840546, 0.7248737812042236, 0.5454342365264893, 0.2521210312843323, 0.2614332437515259, 0.3105127811431885, 0.03566299006342888],
    [0.3035010099411011, 0.07670660316944122, 0.07924123853445053, 0.161861851811409, 0.14367112517356873, 0.06336789578199387, 0.9437791109085083, 0.5998468399047852, 0.4222017228603363, 0.14000535011291504, 0.12471750378608704, 0.31671205163002014, 0.6216381192207336, 0.4062456488609314, 0.1768452525138855, 0.2160402536392212, 0.9336262345314026, 0.289279043674469, ...],
    ...
  ]
>

Nx.take/3 takes and concatenates slices from indices in a given tensor. In this example, each value in tokens represents an index in the first dimension of the weights tensor. What you end up with is a sequence of dense vectors from an original sparse representation. Another option is Nx.take_along_axis/3 which can be used to take indices along an axis of a given tensor. For example, you can combine Nx.take_along_axis/3 with Nx.argsort/2 to sort values along an axis:

a = Nx.tensor([[2, 3, 0, 1, 4, 8], [5, 1, 2, 3, 6, 9]])
indices = Nx.argsort(a, axis: 1)
Nx.take_along_axis(a, indices, axis: 1)

Returns:

#Nx.Tensor<
  s64[2][6]
  [
    [0, 1, 2, 3, 4, 8],
    [1, 2, 3, 5, 6, 9]
  ]
>

The last indexing operation Nx offers is Nx.gather/3. Nx.gather/3 contains a tensor of indices, where the last axes in the indices tensor represents a single value in the source tensor:

t = Nx.tensor([[1, 2], [3, 4]])
Nx.gather(t, Nx.tensor([[1, 1], [0, 1], [1, 0]]))

Returns:

#Nx.Tensor<
  s64[3]
  [4, 2, 3]
>

Notice how the last dimension in the indices tensor has a size of 2 which matches the rank of the input tensor. The last dimension represents indexes in the source tensor: [1, 1], [0, 1], and [1, 0]. The output shape of a gather operation will be equal to the leading dimensions in the indices tensor minus the last dimension.