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, необязательно) – Длина каждого сегмента. По умолчанию: None.
• offsets (Tensor, необязательно) – Смещение каждого сегмента. По умолчанию: None.
• axis (int, необязательно) – Ось, вдоль которой выполняется редукция. По умолчанию: 0.
• unsafe (bool, необязательно) – Пропустить проверку, если True. По умолчанию: False.
• initial (Number, необязательно) – Начальное значение для операции редукции. По умолчанию: 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')