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_elementwise_fn_wrapper

def gather_elementwise_fn_wrapper[dtype: DType, indices_type: DType, InputFnType: def[width: Int, rank: Int, element_alignment: Int](IndexList[rank]) -> SIMD[dtype, width] & ImplicitlyCopyable & RegisterPassable, IndicesFnType: def[width: Int, rank: Int](IndexList[rank]) -> SIMD[indices_type, width] & ImplicitlyCopyable & RegisterPassable, OutputFnType: def[width: SIMDLength, rank: Int, element_alignment: Int](IndexList[rank], SIMD[dtype, width]) -> None & ImplicitlyCopyable & RegisterPassable, *, simd_width: Int, prefetch_fn: OptionalReg[def[input_rank: Int, indices_rank: Int](IndexList[input_rank], IndexList[indices_rank]) capturing thin -> None] = None, target: StringSlice[ImmStaticOrigin] = StringSlice("cpu"), element_alignment: Int = Int(1)](input_fn: InputFnType, indices_fn: IndicesFnType, output_fn: OutputFnType, axis: Axis, input_shape: IndexList[element_type=input_shape.element_type], indices_shape: IndexList[element_type=indices_shape.element_type], output_shape: IndexList[element_type=output_shape.element_type], coords: IndexList[element_type=coords.element_type], error_index_ptr: OptionalReg[Pointer[Int, MutAnyOrigin, _safe=False]] = None) where (eq InputFnType.dtype, dtype) where (eq IndicesFnType.indices_type, indices_type) where (eq OutputFnType.dtype, dtype)

Performs a single elementwise gather step for one output coordinate.

Loads the gather index from indices at the coordinate derived from coords, normalizes it against the input axis size, reads the corresponding value from input, and stores it into output at the original coordinate. On CPU, out-of-bounds indices are recorded through error_index_ptr for later reporting; on GPU, a debug assert traps instead.

Parameters:

dtype (DType): Element type of the input and output values.
indices_type (DType): Element type of the indices buffer.
InputFnType (def[width: Int, rank: Int, element_alignment: Int](IndexList[rank]) -> SIMD[dtype, width] & ImplicitlyCopyable & RegisterPassable): Function type that loads a SIMD vector from the input tensor at given coordinates.
IndicesFnType (def[width: Int, rank: Int](IndexList[rank]) -> SIMD[indices_type, width] & ImplicitlyCopyable & RegisterPassable): Function type that loads a SIMD vector of indices from the indices buffer at given coordinates.
OutputFnType (def[width: SIMDLength, rank: Int, element_alignment: Int](IndexList[rank], SIMD[dtype, width]) -> None & ImplicitlyCopyable & RegisterPassable): Function type that stores a SIMD vector into the output tensor at given coordinates.
simd_width (Int): SIMD vector width for the elementwise load and store.
prefetch_fn (OptionalReg[def[input_rank: Int, indices_rank: Int](IndexList[input_rank], IndexList[indices_rank]) capturing thin -> None]): Optional prefetch callback for software index prefetching (defaults to None).
target (StringSlice[ImmStaticOrigin]): Target backend to execute on, such as "cpu" or "cuda" (defaults to "cpu").
element_alignment (Int): Alignment factor for element loads and stores (defaults to 1).

Args:

input_fn (InputFnType): Callback that loads values from the input tensor.
indices_fn (IndicesFnType): Callback that loads indices from the indices tensor.
output_fn (OutputFnType): Callback that stores values into the output tensor.
axis (Axis): Validated axis along which to gather from the input tensor.
input_shape (IndexList[element_type=input_shape.element_type]): Shape of the input tensor.
indices_shape (IndexList[element_type=indices_shape.element_type]): Shape of the indices tensor.
output_shape (IndexList[element_type=output_shape.element_type]): Shape of the output tensor.
coords (IndexList[element_type=coords.element_type]): Output coordinate being processed in this gather step.
error_index_ptr (OptionalReg[Pointer[Int, MutAnyOrigin, _safe=False]]): Optional pointer to record an out-of-bounds index for CPU error reporting (defaults to None).