torch_gather {torch} | R Documentation |
Gather
Description
Gather
Usage
torch_gather(self, dim, index, sparse_grad = FALSE)
Arguments
self |
(Tensor) the source tensor |
dim |
(int) the axis along which to index |
index |
(LongTensor) the indices of elements to gather |
sparse_grad |
(bool,optional) If |
gather(input, dim, index, sparse_grad=FALSE) -> Tensor
Gathers values along an axis specified by dim
.
For a 3-D tensor the output is specified by::
out[i][j][k] = input[index[i][j][k]][j][k] # if dim == 0 out[i][j][k] = input[i][index[i][j][k]][k] # if dim == 1 out[i][j][k] = input[i][j][index[i][j][k]] # if dim == 2
If input
is an n-dimensional tensor with size
and
dim = i
, then index
must be an -dimensional tensor with
size
where
and
out
will have the same size as index
.
Examples
if (torch_is_installed()) {
t = torch_tensor(matrix(c(1,2,3,4), ncol = 2, byrow = TRUE))
torch_gather(t, 2, torch_tensor(matrix(c(1,1,2,1), ncol = 2, byrow=TRUE), dtype = torch_int64()))
}
[Package torch version 0.13.0 Index]