Skip to main content

ClikaRT::ops::gather

function

gather()​

Tensor gather(
    Tensor x,
    int64_t dim,
    Tensor index
)

Axis-wise gather: read x at positions given by index along dim.

out[i][j][k] = x[index[i][j][k]][j][k] for dim = 0 (likewise for any other dim; only that axis's coordinate is replaced). The output takes index's shape and x's dtype. Index tensors are Int32 or Int64 on every backend; entries must lie in [0, x's dim extent).

Parameters

  • x: the source.
  • dim: the axis the indices address (negative wraps).
  • index: integer positions, same rank as x.

Throws

  • ClikaRT::Error: (INVALID_ARGUMENT) when dim is out of range or the index dtype is not Int32/Int64 (code_name() carries the reason).

Returns: a new tensor shaped like index, dtype of x.

// one logit per row: [B, V] gathered by [B, 1] ids -> [B, 1]
auto picked = ClikaRT::ops::gather(logits, 1, ids);

Declared in ClikaRT/compute/ops.h, line 2407