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

Mojo function

index_tensor

def index_tensor[dtype: DType, indices_type: DType, batch_dims: Int, target: StringSpan[ImmStaticOrigin] = StringSpan("cpu")](data: TileTensor[dtype, Storage=data.Storage, address_space=data.address_space, linear_idx_type=data.linear_idx_type], indices: TileTensor[indices_type, Storage=indices.Storage, address_space=indices.address_space, linear_idx_type=indices.linear_idx_type], output: TileTensor[dtype, Storage=output.Storage, address_space=output.address_space, linear_idx_type=output.linear_idx_type], ctx: DeviceContext)

Index_tensor operation; based on modified implementation of gather_nd.

Parameters:

• ​dtype (DType): Type of data tensor.
• ​indices_type (DType): Type of indices tensor.
• ​batch_dims (Int): Number of batch dimensions. The gather of indexing starts from dimension of data[batch_dims:].
• ​target (StringSpan[ImmStaticOrigin]): The target architecture to execute on.

Args:

• ​data (TileTensor[dtype, Storage=data.Storage, address_space=data.address_space, linear_idx_type=data.linear_idx_type]): Tensor of rank data_rank >= 1.
• ​indices (TileTensor[indices_type, Storage=indices.Storage, address_space=indices.address_space, linear_idx_type=indices.linear_idx_type]): Tensor of rank indices_rank >= 1. All index values are expected to be within bounds [-s, s-1] along axis of size s. It is an error if any of the index values are out of bounds.
• ​output (TileTensor[dtype, Storage=output.Storage, address_space=output.address_space, linear_idx_type=output.linear_idx_type]): Tensor of rank data_rank + indices_rank - indices_shape[-1] - 1 - b.
• ​ctx (DeviceContext): The device context as prepared by the graph compiler.