Function: gather()

gather(src, indices, dst, axis?): Tensor

Defined in: extensions/math.ts:104

Reads one value per lane out of a float32 source tensor, at the positions given by an int32 index tensor. Pairs with argmax, whose output has exactly the shape this expects, so argmax then gather yields the maximum values alongside their indices.

Parameters

src

Tensor

The input float32 source tensor. Shape [d1,...,dk,...,dn].

indices

Tensor

The int32 index tensor, one index per lane. Shape [d1,...,1,...,dn].

dst

Tensor

The pre-allocated float32 destination tensor. Same shape as indices.

axis?

number = -1

The dimension the indices point into. Defaults to -1 (last dimension).

Returns

Tensor

The destination tensor containing the gathered values.

Throws

With code INVALID_ARGUMENT if an index falls outside the gathered axis.