reduce_sum

Description

Reduces dimensions by computing the sum of elements across dimensions.

Prototype

reduce_sum(raw_tensor, axis, keepdims=False)

Parameters

  • raw_tensor: a tvm.tensor for the input tensor

    Atlas training products: supports float16 and float32.

    Atlas inference products: supports float16, float32, and int32.

    Atlas 200I/500 A2 inference products: supports float16 and float32.

    Atlas A2 training products/Atlas A2 inference products: supports float16 and float32.

    Atlas A3 training products/Atlas A3 inference products: supports float16 and float32.

  • axis: an int or a list of ints to reduce. The value range is [–d, d – 1], where d indicates the dimension count of raw_tensor.
  • keepdims: defaults to False, indicating that the axis length is 0 after reduction. For example, if the original shape is (10, 10, 10) and keepdims is set to False, the shape after reduction is (10, 10). If this parameter is set to True, the axis length is 1 after reduction. For example, if the original shape is (10, 10, 10) and keepdims is set to True, the shape after reduction is (10, 10, 1).

Returns

res_tensor: a tvm.tensor for the result tensor

Restrictions

Due to the format restrictions of the computing platform, the reduced data needs to be rearranged for subsequent operations. Do not add vector computations following a reduce operation.

Applicable Models

Atlas training products

Atlas inference products

Atlas 200I/500 A2 inference products

Atlas A2 training products/Atlas A2 inference products

Atlas A3 training products/Atlas A3 inference products

Example

from tbe import tvm
from tbe import dsl
shape = (1024,1024)
input_dtype = "float16"
data = tvm.placeholder(shape, name="data", dtype=input_dtype)
res = dsl.reduce_sum(data, axis=1)