dgl.ops.segment_reduce

dgl.ops.segment_reduce(seglen, value, reducer='sum')

分段归约操作符。

它通过段沿第一维度聚合值张量。 第一个参数 seglen 存储每个段的长度。其 总和必须等于 value 张量的第一维度。 允许零长度的段。

Parameters:

• seglen (Tensor) – 段长度。

• value (Tensor) – 要聚合的值。

• reducer (str, optional) – 聚合方法。可以是 'sum', 'max', 'min', 'mean'。

Returns:

聚合的张量形状为 (len(seglen), value.shape[1:])。

Return type:

张量

示例

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