Mojo module
gather_scatter
comptime values
error_index_fn_type
comptime error_index_fn_type = def(Int) capturing -> None
Structs
-
Axis: -
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
-
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: -
gather_guards: -
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 agatheroperation, 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 agatheroperation, 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 ascatter_elementsoperation, 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 ascatter_ndoperation, and assert the inputs are compatible. -
scatter_set_constant: Scatter the fill_value into the data at the specified indices.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!