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: StringSlice[ImmStaticOrigin] = StringSlice("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 (
StringSlice[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.