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
- Return type
- 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)