triton.language.sum

triton.language.sum(input, axis=None, keep_dims=False, dtype: constexpr | None = None)

返回 input 张量中沿给定 axis 的所有元素的和。

归约操作应具有结合律和交换律。

参数:

• input (Tensor) – 输入值

• axis (int) – 执行归约操作的维度。如果为 None，则归约所有维度

• keep_dims (bool) – 如果为 true，则保留长度为 1 的归约维度

• dtype (tl.dtype) – the desired data type of the returned tensor. If specified, the input tensor is casted to dtype before the operation is performed. This is useful for preventing data overflows. If not specified, integer and bool dtypes are upcasted to tl.int32 and float dtypes are upcasted to at least tl.float32.

This function can also be called as a member function on tensor, as x.sum(...) instead of sum(x, ...).