dgl.ops.segment_reduce

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

分段归约算子。

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

参数:

• seglen (Tensor) – 段长度。

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

• reducer (str, optional) – 聚合方法。可以是 'sum'（求和）、'max'（最大值）、'min'（最小值）或 'mean'（平均值）。

返回:

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

返回类型:

张量

示例

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