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 module

gather_scatter

Implements gather and scatter operations for CPU and GPU, including indexed reductions.

Structs

• ​Axis: Wraps a tensor axis index, optionally normalizing negative values against the tensor rank.
• ​ScatterOobIndexStrategy: Valid indices are within the range [-dim_size, dim_size). Indices which fall outside of that can be handled using different strategies. Note that negative indices are allowed in order to support negative relative indexing. Eg: x[-1] == x[dim_size - 1].

Functions

• ​apply_packed_bitmask: Apply a packed-int32 grammar bitmask to logits in a single fused pass.
• ​gather: Gather operation as defined in https://github.com/onnx/onnx/blob/main/docs/Operators.md#Gather.
• ​gather_elements: Implements ONNX GatherElements op which is equivalent to Pytorch gather.
• ​gather_elementwise_fn_wrapper: Performs a single elementwise gather step for one output coordinate.
• ​gather_guards: Validates that the input, indices, and output shapes are compatible for a gather operation.
• ​gather_nd: GatherND operation as defined in https://github.com/onnx/onnx/blob/main/docs/Operators.md#GatherND. Based on reference implementation: https://github.com/onnx/onnx/blob/main/onnx/backend/test/case/node/gathernd.py.
• ​gather_nd_shape: Compute the output shape of a gather operation, and assert the inputs are compatible.
• ​gather_reduce: Computes output[i, j, k] = input[indices[i, j], k] and simultaneously reduces the output across axis 1 to produce output[i, k].
• ​gather_shape: Compute the output shape of a gather operation, and assert the inputs are compatible.
• ​normalize_neg_index: Indices passed to gather and scatter ops may be negative. This performs a normalization so that they can be used to index into a buffer.
• ​scatter_elements: Implements ONNX ScatterElements op which is equivalent to Pytorch scatter.
• ​scatter_elements_shape: Compute the output shape of a scatter_elements operation, and assert the inputs are compatible.
• ​scatter_nd: Scatter_nd operation without any reduction.
• ​scatter_nd_generator: Implements ONNX ScatterND operation as defined in https://github.com/onnx/onnx/blob/main/docs/Operators.md#ScatterND.
• ​scatter_nd_shape: Compute the output shape of a scatter_nd operation, and assert the inputs are compatible.
• ​scatter_set_constant: Scatter the fill_value into the data at the specified indices.