dgl.ops.segment_reduce

dgl.ops.segment_reduce(seglen, value, reducer='sum')[source]

Segment reduction operator.

It aggregates the value tensor along the first dimension by segments. The first argument seglen stores the length of each segment. Its summation must be equal to the first dimension of the value tensor. Zero-length segments are allowed.

Parameters:

  • seglen (Tensor) – Segment lengths.

  • value (Tensor) – Value to aggregate.

  • reducer (str, optional) – Aggregation method. Can be β€˜sum’, β€˜max’, β€˜min’, β€˜mean’.

Returns:

Aggregated tensor of shape (len(seglen), value.shape[1:]).

Return type:

Tensor

Examples

>>> import dgl
>>> import torch as th
>>> val = th.ones(10, 3)
>>> seg = th.tensor([1, 0, 5, 4])  # 4 segments
>>> dgl.segment_reduce(seg, val)
tensor([[1., 1., 1.],
        [0., 0., 0.],
        [5., 5., 5.],
        [4., 4., 4.]])