Mojo function
sum
sum[axis: Int](inp: LayoutTensor[dtype, layout, origin, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment], out: LayoutTensor[dtype, layout, origin, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=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 = InlineArray[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))
from layout import LayoutTensor, Layout
from layout.math import sum
data = InlineArray[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:
0 1 2
3 4 5
-----
3 5 7
0 1 2
3 4 5
-----
3 5 7
.
Constraints:
All tensors must have statically known shapes.
out.rank must equal inp.rank - 1.
Non-reduction dimensions must match between inp and out.
Currently only supports rank-2 inputs.
Parameters:
- axis (
Int): The axis to sum along.
Args:
- inp (
LayoutTensor[dtype, layout, origin, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment]): The input tensor to sum. - out (
LayoutTensor[dtype, layout, origin, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment]): The output tensor to store sum results.
sum[axis: Int](inp: LayoutTensor[dtype, layout, origin, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment]) -> LayoutTensor[dtype, _reduce_res_row_major_shape(axis, layout), MutableAnyOrigin, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth]
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[dtype, layout, origin, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment]): The input tensor to sum.
Returns:
A new tensor containing the sum values along the specified axis.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!