Mojo function

sum

sum(src: NDBuffer[dtype, 1, origin]) -> SIMD[dtype, 1]

Computes the sum of buffer elements.

Args:

src (NDBuffer[dtype, 1, origin]): The buffer.

Returns:

The sum of the buffer elements.

sum[reduce_axis: Int](src: NDBuffer[dtype, rank, origin, shape, strides, alignment=alignment, address_space=address_space, exclusive=exclusive], dst: NDBuffer[dtype, rank, origin, shape])

Computes the sum across reduce_axis of an NDBuffer.

Parameters:

reduce_axis (Int): The axis to reduce across.

Args:

src (NDBuffer[dtype, rank, origin, shape, strides, alignment=alignment, address_space=address_space, exclusive=exclusive]): The input buffer.
dst (NDBuffer[dtype, rank, origin, shape]): The output buffer.

sum[: origin.set, : origin.set, //, dtype: DType, input_fn: fn[Int, Int](IndexList[$1]) capturing -> SIMD[dtype, $0], output_fn: fn[Int, Int](IndexList[$1], SIMD[dtype, $0]) capturing -> None, /, single_thread_blocking_override: Bool = False, target: StringSlice[StaticConstantOrigin] = __init__[__mlir_type.!kgen.string]("cpu")](input_shape: IndexList[size], reduce_dim: Int, context: DeviceContextPtr = DeviceContextPtr())

Computes the sum across the input and output shape.

This performs the sum computation on the domain specified by input_shape, loading the inputs using the input_fn. The results are stored using the output_fn.

Parameters:

dtype (DType): The dtype of the input and output.
input_fn (fn[Int, Int](IndexList[$1]) capturing -> SIMD[dtype, $0]): The function to load the input.
output_fn (fn[Int, Int](IndexList[$1], SIMD[dtype, $0]) capturing -> None): The function to store the output.
single_thread_blocking_override (Bool): If True, then the operation is run synchronously using a single thread.
target (StringSlice[StaticConstantOrigin]): The target to run on.

Args:

input_shape (IndexList[size]): The input shape.
reduce_dim (Int): The axis to perform the sum on.
context (DeviceContextPtr): The pointer to DeviceContext.

sum[dtype: DType, input_fn_1d: fn[DType, Int](idx: Int) capturing -> SIMD[$0, $1]](length: Int) -> SIMD[dtype, 1]

Computes the sum of a 1D array using a provided input function.

This function performs a reduction (sum) over a 1-dimensional array of the specified length and data type. The input values are provided by the input_fn_1d function, which takes an index and returns a SIMD vector of the specified width and data type. The reduction is performed using a single thread for deterministic results.

Parameters:

dtype (DType): The data type of the elements to sum.
input_fn_1d (fn[DType, Int](idx: Int) capturing -> SIMD[$0, $1]): A function that takes a data type, SIMD width, and index, and returns a SIMD vector of input values.

Args:

length (Int): The number of elements in the 1D array.

Returns:

The sum of all elements as a scalar of the specified data type.

Raises:

Any exception raised by the input function or reduction process.