Function segment_mean

flowdock/utils/model_utils.py:129–146

Computes the mean value of each segment in a tensor.

(src, dst_idx, dst_size)

def segment_mean(src, dst_idx, dst_size):
    """Computes the mean value of each segment in a tensor."""
    out = torch.zeros(
        dst_size,
        *src.shape[1:],
        dtype=src.dtype,
        device=src.device,
    ).index_add_(0, dst_idx, src)
    denom = (
        torch.zeros(
            dst_size,
            *src.shape[1:],
            dtype=src.dtype,
            device=src.device,
        ).index_add_(0, dst_idx, torch.ones_like(src))
        + 1e-8
    )
    return out / denom


def segment_argmin(scores, dst_idx, dst_size, randomize: bool = False) -> torch.Tensor: