tensorplay.functional.take_along_dim

tensorplay.functional.take_along_dim(input, indices, dim=None, *, out=None) → Tensor[source]

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

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

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

Note

This function is similar to NumPy’s take_along_axis. See also tensorplay.gather().

Parameters:

• input (Tensor) – the input tensor.

• indices (LongTensor) – the indices into input. Must have long dtype.

• dim (int, optional) – dimension to select along. Default: 0

Keyword Arguments:

out (Tensor, optional) – the output tensor.

Example:

>>> t = tensorplay.tensor([[10, 30, 20], [60, 40, 50]])
>>> max_idx = tensorplay.argmax(t)
>>> tensorplay.take_along_dim(t, max_idx)
tensor([60])
>>> sorted_idx = tensorplay.argsort(t, dim=1)
>>> tensorplay.take_along_dim(t, sorted_idx, dim=1)
tensor([[10, 20, 30],
        [40, 50, 60]])