megengine.functional.sum

sum(inp, axis=None, keepdims=False)[source]

Calculates the sum of tensor elements over a given axis (or axes).

Parameters
  • inp (Tensor) – input tensor. Should have a numeric data type.

  • axis (Union[int, Sequence[int], None]) – axis or axes along which sums must be computed. By default, the sum must be computed over the entire tensor. If a sequence of integers, sums must be computed over multiple axes.

  • keepdims (bool) – if True, the reduced axes (dimensions) must be included in the result as singleton dimensions, and, accordingly, the result must be compatible with the input tensor (see Broadcasting mechanism and rules). Otherwise, if False, the reduced axes (dimensions) must not be included in the result.

Return type

Tensor

Returns

if the sum was computed over the entire tensor, a zero-dimensional tensor containing the sum; otherwise, a tensor containing the sums. The returned tensor must have a data type determined by Type promotion rules.

Special Cases

Let N equal the number of elements over which to compute the sum.

  • If N is 0, the sum is 0 (i.e., the empty sum).

  • If \(x_i\) is NaN, the sum is NaN (i.e., NaN values propagate).

Warning

If the accumulator is too small, overflow occurs:

>>> x = F.ones(128, dtype="int8")
>>> F.sum(x)
Tensor(-128, dtype=int8, device=xpux:0)

Examples

The sum of an empty tensor is the neutral element 0:

>>> F.sum(Tensor([]))
Tensor(0.0, device=xpux:0)

Normal case:

>>> F.sum(Tensor([1, 2, 3]))
Tensor(6, dtype=int32, device=xpux:0)
>>> F.sum(Tensor([0.5, 1.5]))
Tensor(2.0, device=xpux:0)

Along an axis:

>>> F.sum(Tensor([[1, 2, 3], [4, 5, 6]]), axis=0)
Tensor([5 7 9], dtype=int32, device=xpux:0)
>>> F.sum(Tensor([[1, 2, 3], [4, 5, 6]]), axis=1)
Tensor([ 6 15], dtype=int32, device=xpux:0)