torch.segment_reduce

torch.segment_reduce(data: Tensor, reduce: str, *, lengths: Tensor | None = None, indices: Tensor | None = None, offsets: Tensor | None = None, axis: _int = 0, unsafe: _bool = False, initial: Number | _complex | None = None) → Tensor

在指定轴上对输入张量执行分段归约操作。

参数

• data (Tensor) – 将执行分段归约操作的输入张量。

• reduce (str) – 归约操作的类型。支持的值包括 sum、mean、max、min、prod。

关键字参数

• lengths (Tensor, optional) – 每个段的长度。默认为 None。

• offsets (Tensor, optional) – 每个段的偏移量。默认为 None。

• axis (int, optional) – 用于执行归约的轴。默认为 0。

• unsafe (bool, optional) – 如果为 True，则跳过验证。默认为 False。

• initial (Number, optional) – 归约操作的初始值。默认为 None。

示例

>>> data = torch.tensor([[1, 2, 3, 4],[5, 6, 7, 8],[9, 10, 11, 12]], dtype=torch.float32, device='cuda')
>>> lengths = torch.tensor([2, 1], device='cuda')
>>> torch.segment_reduce(data, 'max', lengths=lengths)
tensor([[ 5.,  6.,  7.,  8.],
        [ 9., 10., 11., 12.]], device='cuda:0')