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

advanced_indexing_getitem

def advanced_indexing_getitem[input_rank: Int, index_rank: Int, input_type: DType, //, start_axis: Int, num_index_tensors: Int, target: StringSlice[ImmStaticOrigin], trace_description: StringSlice[ImmStaticOrigin], InputTensorFn: def[dtype: DType, width: Int](IndexList[input_rank]) -> SIMD[dtype, width] & ImplicitlyCopyable & RegisterPassable, IndicesFn: def[indices_index: Int](IndexList[index_rank]) -> Int & ImplicitlyCopyable & RegisterPassable](out_tensor: TileTensor[input_type, Storage=out_tensor.Storage, address_space=out_tensor.address_space, linear_idx_type=out_tensor.linear_idx_type], in_tensor_strides: IndexList[input_rank], ctx: DeviceContext, input_tensor_fn: InputTensorFn, indices_fn: IndicesFn) where (eq InputTensorFn.input_rank, input_rank) where (eq IndicesFn.index_rank, index_rank)

Implement basic numpy-style advanced indexing.

This is designed to be fused with other view-producing operations to implement full numpy-indexing semantics.

This assumes the dimensions in input_tensor not indexed by index tensors are ":", ie selecting all indices along the slice. For example in numpy:

# rank(indices1) == 3
# rank(indices2) == 3
out_tensor = input_tensor[:, :, :, indices1, indices2, :, :]

We calculate the following for all valid valued indexing variables:

out_tensor[a, b, c, i, j, k, d, e] = input_tensor[
    a, b, c,
    indices1[i, j, k],
    indices2[i, j, k],
    d, e
]

In this example start_axis = 3 and num_index_tensors = 2.

Note: Currently supports contiguous indexing tensors only; boolean tensor masks and view-fusion are not yet implemented.

Parameters:

input_rank (Int): The rank of the input tensor.
index_rank (Int): The rank of the indexing tensors.
input_type (DType): The dtype of the input tensor.
start_axis (Int): The first dimension in input where the indexing tensors are applied. It is assumed the indexing tensors are applied in consecutive dimensions.
num_index_tensors (Int): The number of indexing tensors.
target (StringSlice[ImmStaticOrigin]): The target architecture to operation on.
trace_description (StringSlice[ImmStaticOrigin]): For profiling, the trace name the operation will appear under.
InputTensorFn (def[dtype: DType, width: Int](IndexList[input_rank]) -> SIMD[dtype, width] & ImplicitlyCopyable & RegisterPassable): The type of the input-tensor fusion lambda.
IndicesFn (def[indices_index: Int](IndexList[index_rank]) -> Int & ImplicitlyCopyable & RegisterPassable): The type of the indices fusion lambda.

Args:

out_tensor (TileTensor[input_type, Storage=out_tensor.Storage, address_space=out_tensor.address_space, linear_idx_type=out_tensor.linear_idx_type]): The output tensor to write to.
in_tensor_strides (IndexList[input_rank]): The strides of the input tensor.
ctx (DeviceContext): The device context as prepared by the graph compiler.
input_tensor_fn (InputTensorFn): Fusion lambda for the input tensor.
indices_fn (IndicesFn): Fusion lambda for the indices tensors.