Summing over atoms

metatrain.utils.sum_over_atoms.sum_over_atoms(tensor_map: TensorMap)

A faster version of metatensor.torch.sum_over_samples, specialized for summing over atoms in graph-like TensorMaps.

Parameters:

tensor_map (TensorMap) – The TensorMap to sum over.

Returns:

A new TensorMap with the same keys, but with the samples summed over the atoms.