mindspore.ops.UnsortedSegmentSum

class mindspore.ops.UnsortedSegmentSum[source]

Computes the sum of a tensor along segments.

Calculates a tensor such that \(\text{output}[i] = \sum_{segment\_ids[j] == i} \text{data}[j, \ldots]\), where \(j\) is a tuple describing the index of element in data. segment_ids selects which elements in data to sum up. Segment_ids does not need to be sorted, and it does not need to cover all values in the entire valid value range.

The following figure shows the calculation process of UnsortedSegmentSum:

../../_images/UnsortedSegmentSum.png

Note

If the segment_id i is absent in the segment_ids, then output[i] will be filled with 0.

If the sum of the given segment_ids \(i\) is empty, then \(\text{output}[i] = 0\). If the given segment_ids is negative, the value will be ignored. ‘num_segments’ must be equal to the number of different segment_ids.

Inputs:

  • input_x (Tensor) - The shape is \((x_1, x_2, ..., x_R)\).

  • segment_ids (Tensor) - Set the shape as \((x_1, x_2, ..., x_N)\), where 0 < N <= R.

  • num_segments (int) - Set \(z\) as num_segments.

Outputs:

Tensor, the shape is \((z, x_{N+1}, ..., x_R)\).

Raises

  • TypeError – If num_segments is not an int.

  • ValueError – If length of shape of segment_ids is less than 1.

Supported Platforms:

Ascend GPU CPU

Examples

>>> input_x = Tensor([1, 2, 3, 4], mindspore.float32)
>>> segment_ids = Tensor([0, 0, 1, 2], mindspore.int32)
>>> num_segments = 4
>>> output = ops.UnsortedSegmentSum()(input_x, segment_ids, num_segments)
>>> print(output)
[3. 3. 4. 0.]
>>> input_x = Tensor([1, 2, 3, 4, 2, 5], mindspore.float32)
>>> segment_ids = Tensor([0, 0, 1, 2, 3, 4], mindspore.int32)
>>> num_segments = 6
>>> output = ops.UnsortedSegmentSum()(input_x, segment_ids, num_segments)
>>> print(output)
[3. 3. 4. 2. 5. 0.]