Mojo function

scatter_elements_shape

scatter_elements_shape[rank: Int, input_type: DType, indices_type: DType, //, *, single_thread_blocking_override: Bool](input: NDBuffer[input_type, rank, origin], updates: NDBuffer[input_type, rank, origin], indices: NDBuffer[indices_type, rank, origin], axis: Int) -> Index[rank]

Compute the output shape of a scatter_elements operation, and assert the inputs are compatible.

Parameters:

rank (Int): Rank of the input tensor.
input_type (DType): Type of the input tensor.
indices_type (DType): Type of the indices tensor.
single_thread_blocking_override (Bool): If True, then the operation is run synchronously using a single thread.

Args:

input (NDBuffer[input_type, rank, origin]): The input tensor.
updates (NDBuffer[input_type, rank, origin]): The input tensor.
indices (NDBuffer[indices_type, rank, origin]): The indices tensor.
axis (Int): The axis.

Returns:

The output shape.