For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

gather_elements

def gather_elements[input_type: DType, indices_type: DType](input: TileTensor[input_type, Storage=input.Storage, address_space=input.address_space, linear_idx_type=input.linear_idx_type], indices: TileTensor[indices_type, Storage=indices.Storage, address_space=indices.address_space, linear_idx_type=indices.linear_idx_type], _axis: Int, output: TileTensor[input_type, Storage=output.Storage, address_space=output.address_space, linear_idx_type=output.linear_idx_type], ctx: DeviceContext)

Implements ONNX GatherElements op which is equivalent to Pytorch gather.

Parameters:

input_type (DType): Element type of input and output.
indices_type (DType): Element type of the indices tensor, must be int32 or int64.

Args:

input (TileTensor[input_type, Storage=input.Storage, address_space=input.address_space, linear_idx_type=input.linear_idx_type]): Source tensor to gather values from.
indices (TileTensor[indices_type, Storage=indices.Storage, address_space=indices.address_space, linear_idx_type=indices.linear_idx_type]): Tensor of indices to gather along _axis, shaped the same as output. Values must be in range [-dim, dim) along _axis.
_axis (Int): Axis along which to gather from input, in range [-rank, rank).
output (TileTensor[input_type, Storage=output.Storage, address_space=output.address_space, linear_idx_type=output.linear_idx_type]): Destination tensor for gathered values, shaped the same as indices.
ctx (DeviceContext): Device context for execution.