gather

→Tensor

gather(input: Tensor, dim: _int, index: Tensor)

source edit

Implementing kernel

C++gather_opfree fn

Gather values along dim according to index.

For each output position the value is read from input with index substituted for the dim-th coordinate. All other coordinates are inherited from the output position itself — a "row-wise" lookup. Tensor.gather(dim, index) is the same op.

Parameters

inputTensor

Source tensor.

dimint

Axis along which to gather. Negative values count from the end.

indexTensor

Integer tensor with as many dimensions as input, no larger than it along any axis but dim. Values must be in [0, input.size(dim)).

Returns

Tensor

Tensor with the same shape as index.

Notes

Mathematically,

\text{out}[i_0, \dots, i_d, \dots] \;=\; \text{input}[i_0, \dots, \text{index}[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, 1, idx)
tensor([[1, 3], [5, 5]], dtype=lucid.int64)