Fix aggregate nodes when last batch in a graph has no edges/nodes (#119)

Author: vsimkus

orb_models/forcefield/segment_ops.py

@@ -44,11 +44,11 @@ def aggregate_nodes(
         torch.use_deterministic_algorithms(True)
     segments = torch.arange(count, device=device).repeat_interleave(n_node)
     if reduction == "sum":
-        return scatter_sum(tensor, segments, dim=0)
+        return scatter_sum(tensor, segments, dim=0, dim_size=count)
     elif reduction == "mean":
-        return scatter_mean(tensor, segments, dim=0)
+        return scatter_mean(tensor, segments, dim=0, dim_size=count)
     elif reduction == "max":
-        return segment_max(tensor, segments, count)
+        return segment_max(tensor, segments, num_segments=count)
     else:
         raise ValueError("Invalid reduction argument. Use sum, mean or max.")
 

tests/atomgraphs/test_segment_ops.py

@@ -58,6 +58,16 @@ def test_aggregate_nodes(reduction, dtype):
     assert torch.allclose(res[2, :], reduce_fn(tensor[8:, :]))
 
 
+@pytest.mark.parametrize("reduction", ["sum", "mean", "max"])
+def test_aggregate_nodes_with_zero_size_last_graph(reduction):
+    tensor = torch.randn(10, 10)
+    n_node = torch.tensor([8, 2, 0, 0], dtype=torch.int32)
+
+    res = segment_ops.aggregate_nodes(tensor, n_node=n_node, reduction=reduction)
+
+    assert res.shape == (4, 10)
+
+
 @pytest.mark.parametrize(
     "dtype", [torch.float32, torch.float64, torch.int32, torch.int64]
 )