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