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

gather_guards

def gather_guards(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])

Validates that the input, indices, and output shapes are compatible for a gather operation.

Args:

• ​axis (Axis): Axis along which the gather is performed; must be non-negative and less than input_shape rank.
• ​input_shape (IndexList[element_type=input_shape.element_type]): Shape of the input tensor being gathered from.
• ​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; must match the gather output shape derived from input_shape and indices_shape.

Raises:

If the axis is negative, out of range, or the output shape does not match the expected gather shape derived from the input and indices.