take_along_dim(input: Tensor, indices: Tensor, dim: _int | None = ...)Gather elements from input at positions indices along dim.
Advanced indexing primitive analogous to np.take_along_axis:
selects one element from input for every entry in indices,
broadcasting the remaining (non-dim) axes between the two
tensors — lucid.gather after that broadcast, which gather
itself does not do. With dim=None both are read flattened.
Parameters
inputTensorSource tensor.
indicesTensorInteger tensor of positions along
dim, with as many dims as
input. Its shape must be broadcast-compatible with input on every
axis other than dim; the size along dim controls the size
of the output along that axis. A negative position counts from
the end.dimintAxis along which to gather. Negative values count from the end.
None (default) reads input and indices flattened.Returns
TensorTensor with the broadcast shape of input and indices (with
the size along dim taken from indices), 1-D when dim
is None. Dtype matches input.
Raises
lucid._C.engine.ShapeMismatchindices has another number of dims than input.IndexErrorThe shapes do not broadcast on an axis other than
dim.Notes
For 1-D inputs the operation reduces to plain integer indexing. For higher-rank inputs, the result satisfies
Typical uses include "gather the top-k elements per row" patterns
after lucid.argsort / lucid.topk.
Examples
>>> import lucid
>>> x = lucid.tensor([[1.0, 2.0, 3.0],
... [4.0, 5.0, 6.0]])
>>> idx = lucid.tensor([[2, 0], [1, 2]])
>>> lucid.take_along_dim(x, idx, dim=1)
tensor([[3., 1.], [5., 6.]])