torch (version 0.14.2)

torch_take_along_dim: Selects values from input at the 1-dimensional indices from indices along the given dim.

Description

Selects values from input at the 1-dimensional indices from indices along the given dim.

Usage

torch_take_along_dim(self, indices, dim = NULL)

Arguments

self

the input tensor.

indices

the indices into input. Must have long dtype.

dim

the dimension to select along. Default is NULL.

Examples

if (torch_is_installed()) {
t <- torch_tensor(matrix(c(10, 30, 20, 60, 40, 50), nrow = 2))
max_idx <- torch_argmax(t)
torch_take_along_dim(t, max_idx)

sorted_idx <- torch_argsort(t, dim=2)
torch_take_along_dim(t, sorted_idx, dim=2)

}