Ragged segment sum and mean

Your task

Function signature

segment_sum_mean(
    values: torch.Tensor,
    segment_id: torch.Tensor,
    num_segments: int,
) -> dict[str, torch.Tensor]

Worked example

Repeated and empty segments

import torch

values = torch.tensor([[2.0, 4.0], [6.0, 8.0], [100.0, -100.0], [-2.0, 0.0]], dtype=torch.float64, requires_grad=True)
segment_id = torch.tensor([1, 1, -1, 0], dtype=torch.int64)
num_segments = 4

result = segment_sum_mean(values=values, segment_id=segment_id, num_segments=num_segments)

Expected return

{
    'sums': torch.tensor([[-2.0, 0.0], [8.0, 12.0], [0.0, 0.0], [0.0, 0.0]], dtype=torch.float64),
    'means': torch.tensor([[-2.0, 0.0], [4.0, 6.0], [0.0, 0.0], [0.0, 0.0]], dtype=torch.float64),
    'counts': torch.tensor([1, 2, 0, 0], dtype=torch.int64),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Group values[E,\ast F] by integer segment_id [E] along axis zero only, retaining every trailing feature dimension \ast F. The trailing dimensions may be absent (scalar rows), multiple, or empty. Required num_segments=N is a nonnegative integer; every ID is either -1 (padding) or in [0,N).
• For each s in [0,N), \mathrm{counts}[s]=\mathrm{number} of rows with \mathrm{segment\_id}[e]=s, \mathrm{sums}[s]=\sum_{e:\mathrm{segment\_id}[e]=s} values [e], and \mathrm{means}[s]=\mathrm{sums}[s]/\operatorname{max}(\mathrm{counts}[s],1). Return all three tensors.
• Repeated and unsorted IDs accumulate; padding rows do not affect sums, means, counts or gradients. Missing intermediate and trailing segments must remain in the output with \mathrm{sums}=\mathrm{means}=0 and \mathrm{counts}=0. E=0 is legal; N=0 is legal only when there are no non-padding rows.
• sums and means have shape [N,\ast F] and preserve values' floating dtype/device; counts is int64 [N] on the same device. values and IDs share a device, and all floating values are finite. There is no \epsilon, weighting, sorting requirement, feature-axis reduction, residual addition or input mutation.
• Both contiguous and strided input layouts must work.

Why this matters

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def segment_sum_mean(
    values: torch.Tensor,  # (E,f_1,...,f_k): any trailing feature axes, including none.
    segment_id: torch.Tensor,  # (E,)
    num_segments: int,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    sums: (N,f_1,...,f_k): preserve every trailing axis of values.
    means: (N,f_1,...,f_k): preserve every trailing axis of values.
    counts: (N,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement segment_sum_mean")

Sources and further reading

• vgsatorras/egnn · egnn_clean.py
• chemprop/chemprop · agg.py