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

scatter_elements

def scatter_elements[rank: Int, input_type: DType, indices_type: DType, *, reduce_fn: OptionalReg[def[dtype: DType, width: SIMDLength](SIMD[dtype, width], SIMD[dtype, width]) thin -> SIMD[dtype, width]] = None](input: ManagedTensorSlice[static_spec=input.static_spec], indices: ManagedTensorSlice[static_spec=indices.static_spec], updates: ManagedTensorSlice[static_spec=updates.static_spec], _axis: Int, output: ManagedTensorSlice[static_spec=output.static_spec], ctx: DeviceContext)

Implements ONNX ScatterElements op which is equivalent to Pytorch scatter.

Parameters:

rank (Int): Rank of the input, indices, updates, and output tensors.
input_type (DType): Element type of input, updates, and output.
indices_type (DType): Element type of indices (must be int32 or int64).
reduce_fn (OptionalReg[def[dtype: DType, width: SIMDLength](SIMD[dtype, width], SIMD[dtype, width]) thin -> SIMD[dtype, width]]): Reduction function to apply: none (default, overwrite), add, mul, max, min. Updates for duplicate indices are reduced atomically, in unspecified order (without a reduce_fn, duplicates leave an unspecified winner instead).

Args:

input (ManagedTensorSlice[static_spec=input.static_spec]): Source tensor copied into output before scattering.
indices (ManagedTensorSlice[static_spec=indices.static_spec]): Indices along _axis selecting where updates land in output.
updates (ManagedTensorSlice[static_spec=updates.static_spec]): Values to scatter into output at positions given by indices.
_axis (Int): Axis along which to scatter; must be in [-rank, rank).
output (ManagedTensorSlice[static_spec=output.static_spec]): Output tensor, same shape as input, receiving scattered updates.
ctx (DeviceContext): Device context for execution.