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

In torch: Tensors and Neural Networks with 'GPU' Acceleration

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.

Note

If dim is NULL, the input array is treated as if it has been flattened to 1d.

Functions that return indices along a dimension, like torch_argmax() and torch_argsort(), are designed to work with this function. See the examples below.

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)

}

torch documentation built on Aug. 21, 2025, 5:50 p.m.