take_op

→TensorImplPtr

int take_op(const int & a, const int & indices)

source edit

Python wrappers

pylucid.ops.takefunction

Index a with a flat index list: $y_j = \mathrm{flatten}(a)_{i_j}$ .

Composite over reshape_op + gather_op. The input is flattened to a 1-D view first so the index dimension is unambiguous; gradient flows back through GatherBackward (scatter-add) and ReshapeBackward.

Math

y_j = a_{\mathrm{unravel}(i_j)}

where $\mathrm{unravel}$ maps the linear index into a's row-major coordinates.

See Also

index_select_op — keeps the original rank instead of flattening.

Parameters

aTensorImplPtr

Source tensor of arbitrary shape.

indicesTensorImplPtr

Integer tensor (int32 or int64) of arbitrary shape; each element is a flat-index into a.

Returns

TensorImplPtr

Tensor of the same shape as indices and dtype of a.

Raises

Failure

If either input is null, or if indices is not an integer dtype.

Notes

Out-of-bounds indices are forwarded verbatim to gather_op — behaviour matches the underlying primitive's bounds policy.