nvalchemi.training.losses.reductions.per_graph_sum

nvalchemi.training.losses.reductions.per_graph_sum(values, batch_idx, num_graphs=None)

Sum per-node values into per-graph values via scatter_add_.

See the module docstring for batch_idx / num_graphs semantics and error conditions.

Returns:

Per-graph sums of shape (num_graphs, *values.shape[1:]).

Return type:

Float[torch.Tensor, “B …”]

Parameters:

• values (Float[torch.Tensor, 'V ...'])

• batch_idx (BatchIndices)

• num_graphs (int | None)