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