segment_sum

paddle.incubate. segment_sum ( data, segment_ids, name=None ) [source]

Segment Sum Operator.

Sum the elements of input data which with the same index in segment_ids. It computes a tensor such that

\[out_i = \sum_{j \in \{segment\_ids_j == i \} } data_{j}\]

where sum is over j such that segment_ids[j] == i.

Parameters
  • data (Tensor) – A tensor, available data type float32, float64, int32, int64.

  • segment_ids (Tensor) – A 1-D tensor, which have the same size with the first dimension of input data. Available data type is int32, int64.

  • name (str, optional) – Name for the operation (optional, default is None). For more information, please refer to Name.

Returns

Tensor, the Segment Sum result.

Examples

>>> import paddle
>>> data = paddle.to_tensor([[1, 2, 3], [3, 2, 1], [4, 5, 6]], dtype='float32')
>>> segment_ids = paddle.to_tensor([0, 0, 1], dtype='int32')
>>> out = paddle.incubate.segment_sum(data, segment_ids)
>>> print(out)
Tensor(shape=[2, 3], dtype=float32, place=Place(cpu), stop_gradient=True,
[[4., 4., 4.],
 [4., 5., 6.]])