gather

→Tensor

gather(input: Tensor, indices: Tensor, dim: DimLike = ...)

source edit

Implementing kernel

C++gather_opfree fn

Gather values along dim according to indices.

For each output position the value is read from input with indices substituted for the dim-th coordinate. All other coordinates are inherited from the output position itself — a "row-wise" lookup.

Parameters

inputTensor

Source tensor.

indicesTensor

Integer tensor whose shape matches input except possibly along dim. Values must be in [0, input.size(dim)).

dimDimLike= ...

Axis along which to gather.

Returns

Tensor

Tensor with the same shape as indices.

Notes

Mathematically,

\text{out}[i_0, \dots, i_d, \dots] \;=\; \text{input}[i_0, \dots, \text{indices}[i_0, \dots, i_d, \dots], \dots] .

Inverse op is scatter.

Examples

>>> import lucid
>>> x = lucid.tensor([[1, 2, 3], [4, 5, 6]])
>>> idx = lucid.tensor([[0, 2], [1, 1]])
>>> lucid.gather(x, idx, dim=1)
Tensor([[1, 3],
        [5, 5]])