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

row_reduce

def row_reduce[BLOCK_SIZE: Int, input_fn: def[dtype: DType, width: Int, rank: Int](IndexList[rank]) capturing thin -> SIMD[dtype, width], reduce_fn: def[dtype: DType, width: SIMDLength](SIMD[dtype, width], SIMD[dtype, width]) capturing thin -> SIMD[dtype, width], dtype: DType, simd_width: Int, rank: Int, accum_type: DType = get_accum_type[dtype](), *, axis: Int](mut row_coords: IndexList[rank], init: Scalar[dtype], row_size: Int) -> Scalar[accum_type]

Reduces a single row along the given axis using block-level cooperative reduction. Delegates to the multi-reduction row_reduce overload with num_reductions=1.

Parameters:

BLOCK_SIZE (Int): The number of threads per block.
input_fn (def[dtype: DType, width: Int, rank: Int](IndexList[rank]) capturing thin -> SIMD[dtype, width]): The lambda to load input elements.
reduce_fn (def[dtype: DType, width: SIMDLength](SIMD[dtype, width], SIMD[dtype, width]) capturing thin -> SIMD[dtype, width]): The binary reduction function.
dtype (DType): The data type of the input elements.
simd_width (Int): The SIMD vector width.
rank (Int): The tensor rank.
accum_type (DType): The accumulator data type (defaults to widened type).
axis (Int): The axis along which to reduce.

Args:

row_coords (IndexList[rank]): The ND coordinates identifying the row.
init (Scalar[dtype]): The identity value for the reduction.
row_size (Int): The number of elements in the row.

Returns:

Scalar[accum_type]: The reduced scalar result for the row.

def row_reduce[BLOCK_SIZE: Int, num_reductions: Int, input_fn: def[dtype: DType, width: Int, rank: Int](IndexList[rank]) capturing thin -> SIMD[dtype, width], reduce_fn: def[dtype: DType, width: SIMDLength, reduction_idx: Int](SIMD[dtype, width], SIMD[dtype, width]) capturing thin -> SIMD[dtype, width], dtype: DType, simd_width: Int, rank: Int, accum_type: DType = get_accum_type[dtype](), *, axis: Int](mut row_coords: IndexList[rank], init: StaticTuple[Scalar[dtype], num_reductions], row_size: Int) -> StaticTuple[Scalar[accum_type], num_reductions]

Reduces a row along the given axis with multiple fused reductions using cooperative SIMD-width reads across threads, a block_reduce, and scalar tail handling.

Parameters:

BLOCK_SIZE (Int): The number of threads per block.
num_reductions (Int): The number of fused reductions to perform.
input_fn (def[dtype: DType, width: Int, rank: Int](IndexList[rank]) capturing thin -> SIMD[dtype, width]): The lambda to load input elements.
reduce_fn (def[dtype: DType, width: SIMDLength, reduction_idx: Int](SIMD[dtype, width], SIMD[dtype, width]) capturing thin -> SIMD[dtype, width]): The binary reduction function parameterized by reduction index.
dtype (DType): The data type of the input elements.
simd_width (Int): The SIMD vector width.
rank (Int): The tensor rank.
accum_type (DType): The accumulator data type (defaults to widened type).
axis (Int): The axis along which to reduce.

Args:

row_coords (IndexList[rank]): The ND coordinates identifying the row.
init (StaticTuple[Scalar[dtype], num_reductions]): The identity values for each reduction.
row_size (Int): The number of elements in the row.

Returns:

StaticTuple[Scalar[accum_type], num_reductions]: The reduced scalar results, one per fused reduction.