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.]])