xnn.common.models.ops.scatter_sum#

xnn.common.models.ops.scatter_sum(src, index, dim_size)[source]#

Sum rows of src into dim_size buckets given by index (dim 0).

This is the canonical edge -> node aggregation for message passing, implemented with torch.Tensor.index_add(). For each row e it accumulates out[index[e]] += src[e].

Parameters:

src (torch.Tensor) – Source tensor whose leading dimension is scattered. Shape (E, *feature_dims) (e.g. per-edge messages).
index (torch.Tensor) – 1-D integer tensor of length E giving, for each row of src, the bucket (destination node) it is added into along dimension 0.
dim_size (int) – Number of output buckets, i.e. the size of the leading dimension of the result (e.g. the number of nodes).

Returns:

Tensor of shape (dim_size, *feature_dims) with the same dtype and device as src, holding the per-bucket sums.

Return type:

torch.Tensor

Notes

Written to stay compatible with torch.jit.script() so scriptable model cores (e.g. SchNet’s node_energy) can call it.