PyTorch `segment_sum` operator

segment_sum.py

from typing import Optional
import math

from torch import Tensor
import torch

def segment_sum(input: Tensor, indexes: Tensor, n: Optional[int] = None, **kwargs) -> Tensor:
    if indexes.ndim == 1:
        indexes = torch.repeat_interleave(indexes, math.prod([*input.shape[1:]])).view(*[indexes.shape[0], *input.shape[1:]])

    if n is None:
        n = max([*indexes]) + 1

    return torch.zeros(n, *input.shape[1:]).scatter_add(0, indexes, input.to(torch.float32)).to(**kwargs)

input = torch.tensor([[1, 2, 3, 4], [4, 3, 2, 1], [5, 6, 7, 8]])

indexes = torch.tensor([0, 0, 1])

segment_sum(input, indexes, 2, dtype=torch.int32)

tensor([[5, 5, 5, 5],
        [5, 6, 7, 8]], dtype=torch.int32)