megengine.functional.gather#

gather(inp, axis, index)[source]#

Gathers data from input tensor on axis using index.

For a 3-D tensor, the output is specified by:

out[i][j][k] = inp[index[i][j][k]][j][k] # if axis == 0
out[i][j][k] = inp[i][index[i][j][k]][k] # if axis == 1
out[i][j][k] = inp[i][j][index[i][j][k]] # if axis == 2

if input tensor is a n-dimensional tensor with size \((x_0,x_1,...,x_{i-1},x_i,x_{i+1},...,x_{n-1})\) and axis=i, then index must be a n-dimensional tensor with size \((x_0,x_1,...,x_{i-1},y,x_{i+1},...,x_{n-1})\) where \(y\ge 1\) and output will have the same size as index.

Parameters:

inp (Tensor) – input tensor.
axis (int) – along which axis to index.
index (Tensor) – indices of elements to gather.

Return type:

Tensor

Returns:

output tensor.

Examples

>>> inp = Tensor([
...     [1,2], [3,4], [5,6],
... ])
>>> index = Tensor([[0,2], [1,0]])
>>> F.gather(inp, 0, index)
Tensor([[1 6]
 [3 2]], dtype=int32, device=xpux:0)