Fix a bug in the QuantileAggregation with the dim size parameter passed

CHANGELOG.md

@@ -45,6 +45,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Fixed an index-out-of-range bug in `QuantileAggregation` when `dim_size` is passed ([#7407](https://github.com/pyg-team/pytorch_geometric/pull/7407))
 - The `filter_per_worker` option will not get automatically inferred by default based on the device of the underlying data ([#7399](https://github.com/pyg-team/pytorch_geometric/pull/7399))
 - Fixed a bug in `LightGCN.recommendation_loss()` to only use the embeddings of the nodes involved in the current mini-batch ([#7384](https://github.com/pyg-team/pytorch_geometric/pull/7384))
 - Added an optional `max_num_elements` argument to `SortAggregation` ([#7367](https://github.com/pyg-team/pytorch_geometric/pull/7367))

test/nn/aggr/test_quantile.py

@@ -7,7 +7,9 @@
 @pytest.mark.parametrize('q', [0., .1, .2, .3, .4, .5, .6, .7, .8, .9, 1.])
 @pytest.mark.parametrize('interpolation', QuantileAggregation.interpolations)
 @pytest.mark.parametrize('dim', [0, 1])
-def test_quantile_aggregation(q, interpolation, dim):
+@pytest.mark.parametrize('dim_size', [None, 15])
+@pytest.mark.parametrize('fill_value', [0.0, 10.0])
+def test_quantile_aggregation(q, interpolation, dim, dim_size, fill_value):
     x = torch.tensor([
         [0.0, 1.0, 2.0],
         [3.0, 4.0, 5.0],
@@ -22,12 +24,18 @@ def test_quantile_aggregation(q, interpolation, dim):
     ])
     index = torch.zeros(x.size(dim), dtype=torch.long)
 
-    aggr = QuantileAggregation(q=q, interpolation=interpolation)
+    aggr = QuantileAggregation(q=q, interpolation=interpolation,
+                               fill_value=fill_value)
     assert str(aggr) == f"QuantileAggregation(q={q})"
 
-    out = aggr(x, index, dim=dim)
+    out = aggr(x, index, dim=dim, dim_size=dim_size)
     expected = x.quantile(q, dim, interpolation=interpolation, keepdim=True)
-    assert torch.allclose(out, expected)
+
+    assert torch.allclose(out.narrow(dim, 0, 1), expected)
+
+    if out.size(0) > index.max() + 1:
+        padding = out.narrow(dim, 1, out.size(dim) - 1)
+        assert torch.allclose(padding, torch.tensor(fill_value))
 
 
 def test_median_aggregation():

torch_geometric/nn/aggr/quantile.py

@@ -78,6 +78,11 @@ def forward(self, x: Tensor, index: Optional[Tensor] = None,
         count = torch.bincount(index, minlength=dim_size or 0)
         cumsum = torch.cumsum(count, dim=0) - count
 
+        # In case there exists dangling indices (`dim_size > index.max()`), we
+        # need to clamp them to prevent out-of-bound issues:
+        if dim_size is not None:
+            cumsum = cumsum.clamp(max=x.size(dim) - 1)
+
         q_point = self.q * (count - 1) + cumsum
         q_point = q_point.t().reshape(-1)