mindspore.ops.unsorted_segment_sum

mindspore.ops.unsorted_segment_sum(input_x, segment_ids, num_segments)

分段计算输入tensor的和。

unsorted_segment_sum的计算过程如下图所示：

$$\text{output}[i]=\sum _{segment\mathrm{\_}ids[j]==i}\text{data}[j,\dots ]$$

其中 $j,...$ 是代表元素索引的tuple。 segment_ids 确定输入tensor元素的分段。 segment_ids 不需要排序，也不需要覆盖 num_segments 范围内的所有值。

说明

• 如果 segment_ids 中不存在segment_id i ，则对输出 output[i] 填充0。

• 在Ascend平台上，如果segment_id的值小于0或大于输入tensor的shape的长度，将触发执行错误。

如果给定的segment_ids $i$ 的和为空，则 $\text{output}[i]=0$ 。如果 segment_ids 元素为负数，将忽略该值。 num_segments 必须等于不同segment_id的数量。

参数：

• input_x (Tensor) - 输入tensor。

• segment_ids (Tensor) - 标识每个元素所属段。

• num_segments (Union[int, Tensor], 可选) - 分段数量，可以为int或零维的tensor。

返回：

Tensor

支持平台：

Ascend GPU CPU

样例：

>>> import mindspore
>>> input_x = mindspore.tensor([1, 2, 3, 4])
>>> segment_ids = mindspore.tensor([0, 0, 1, 2])
>>> num_segments = 3
>>> mindspore.ops.unsorted_segment_sum(input_x, segment_ids, num_segments)
Tensor(shape=[3], dtype=Int64, value= [3, 3, 4])
>>> input_x = mindspore.tensor([1, 2, 3, 4, 2, 5])
>>> segment_ids = mindspore.tensor([0, 0, 1, 2, 3, 4])
>>> num_segments = 6
>>> mindspore.ops.unsorted_segment_sum(input_x, segment_ids, num_segments)
Tensor(shape=[6], dtype=Int64, value= [3, 3, 4, 2, 5, 0])