r""" Gathers data from input tensor on axis using index. For a 3-D tensor, the output is specified by: .. code-block:: 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]
(inp: Tensor, axis: int, index: Tensor)
| 680 | |
| 681 | |
| 682 | def gather(inp: Tensor, axis: int, index: Tensor) -> Tensor: |
| 683 | # TODO: rewrite doc |
| 684 | r""" |
| 685 | Gathers data from input tensor on axis using index. |
| 686 | |
| 687 | For a 3-D tensor, the output is specified by: |
| 688 | |
| 689 | .. code-block:: |
| 690 | |
| 691 | out[i][j][k] = inp[index[i][j][k]][j][k] # if axis == 0 |
| 692 | out[i][j][k] = inp[i][index[i][j][k]][k] # if axis == 1 |
| 693 | out[i][j][k] = inp[i][j][index[i][j][k]] # if axis == 2 |
| 694 | |
| 695 | if input tensor is a n-dimensional tensor with size |
| 696 | :math:`(x_0,x_1,...,x_{i-1},x_i,x_{i+1},...,x_{n-1})` and axis=i, |
| 697 | then index must be a n-dimensional tensor with size |
| 698 | :math:`(x_0,x_1,...,x_{i-1},y,x_{i+1},...,x_{n-1})` where :math:`y\ge 1` and |
| 699 | output will have the same size as index. |
| 700 | |
| 701 | Args: |
| 702 | inp: input tensor. |
| 703 | axis: along which axis to index. |
| 704 | index: indices of elements to gather. |
| 705 | |
| 706 | Return: |
| 707 | output tensor. |
| 708 | |
| 709 | Examples: |
| 710 | >>> inp = Tensor([ |
| 711 | ... [1,2], [3,4], [5,6], |
| 712 | ... ]) |
| 713 | >>> index = Tensor([[0,2], [1,0]]) |
| 714 | >>> F.gather(inp, 0, index) |
| 715 | Tensor([[1 6] |
| 716 | [3 2]], dtype=int32, device=xpux:0) |
| 717 | """ |
| 718 | input_shape = inp.shape |
| 719 | index_shape = index.shape |
| 720 | input_dims = len(input_shape) |
| 721 | index_dims = len(index_shape) |
| 722 | if input_dims != index_dims: |
| 723 | raise ValueError( |
| 724 | "The index tensor must have same dimensions as input tensor, " |
| 725 | "But the input dims:{}, the index dims:{}".format(input_dims, index_dims) |
| 726 | ) |
| 727 | |
| 728 | idx = _get_idx(index, axis) |
| 729 | return inp[idx].reshape(index.shape) # pylint: disable=no-member |
| 730 | |
| 731 | |
| 732 | def scatter(inp: Tensor, axis: int, index: Tensor, source: Tensor) -> Tensor: |