For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo function

sum

def sum[axis: Int](inp: LayoutTensor[address_space=inp.address_space, element_layout=inp.element_layout, layout_int_type=inp.layout_int_type, linear_idx_type=inp.linear_idx_type, masked=inp.masked, alignment=inp.alignment], outp: LayoutTensor[address_space=outp.address_space, element_layout=outp.element_layout, layout_int_type=outp.layout_int_type, linear_idx_type=outp.linear_idx_type, masked=outp.masked, alignment=outp.alignment])

Computes sum reduction along specified axis.

Reduces the input tensor by summing elements along the specified axis and stores the result in the output tensor.

Example:

from layout import LayoutTensor, Layout
from layout.math import sum

data: Array[Int32, 6] = [0, 1, 2, 3, 4, 5]
tensor = LayoutTensor[DType.int32, Layout.row_major(2, 3)](data)
print(tensor)
print("-----")
print(sum[0](tensor))

Output:

1 2
4 5
-----
5 7

Constraints:

All tensors must have statically known shapes. outp.rank must equal inp.rank - 1. Non-reduction dimensions must match between inp and outp. Currently only supports rank-2 inputs.

Parameters:

axis (Int): The axis to sum along.

Args:

inp (LayoutTensor[address_space=inp.address_space, element_layout=inp.element_layout, layout_int_type=inp.layout_int_type, linear_idx_type=inp.linear_idx_type, masked=inp.masked, alignment=inp.alignment]): The input tensor to sum.
outp (LayoutTensor[address_space=outp.address_space, element_layout=outp.element_layout, layout_int_type=outp.layout_int_type, linear_idx_type=outp.linear_idx_type, masked=outp.masked, alignment=outp.alignment]): The output tensor to store sum results.

def sum[axis: Int](inp: LayoutTensor[address_space=inp.address_space, element_layout=inp.element_layout, layout_int_type=inp.layout_int_type, linear_idx_type=inp.linear_idx_type, masked=inp.masked, alignment=inp.alignment]) -> LayoutTensor[inp.dtype, _reduce_res_row_major_shape(axis, inp.layout), MutAnyOrigin, address_space=inp.address_space, element_layout=inp.element_layout, layout_int_type=inp.layout_int_type, linear_idx_type=inp.linear_idx_type]

Computes sum reduction along specified axis, returning a new tensor.

Reduces the input tensor by summing elements along the specified axis and returns a new tensor with the results.

Constraints:

All tensors must have statically known shapes. Result will have rank equal to inp.rank - 1. Non-reduction dimensions in the result match the input. Currently only supports rank-2 inputs.

Parameters:

axis (Int): The axis to sum along.

Args:

inp (LayoutTensor[address_space=inp.address_space, element_layout=inp.element_layout, layout_int_type=inp.layout_int_type, linear_idx_type=inp.linear_idx_type, masked=inp.masked, alignment=inp.alignment]): The input tensor to sum.

Returns:

LayoutTensor[inp.dtype, _reduce_res_row_major_shape(axis, inp.layout), MutAnyOrigin, address_space=inp.address_space, element_layout=inp.element_layout, layout_int_type=inp.layout_int_type, linear_idx_type=inp.linear_idx_type]: A new tensor containing the sum values along the specified axis.