Repository: pyg-team/pytorch_geometric

Issue:
Median Aggregation Support
### 🚀 The feature, motivation and pitch

I'm working on a noise graph, some papers show that the interference of noise can be largely reduced using median aggregation.

Papers:
Robustness of Graph Neural Networks at Scale [NIPS'21]
Understanding Structural Vulnerability in Graph Convolutional Networks [IJCAI'21]

### Alternatives

Using torch.argsort like https://github.com/EdisonLeeeee/MedianGCN in the large-scale graph is time costing.
The CUDA-accelerated version of median aggregation is implemented here,https://github.com/sigeisler/robustness_of_gnns_at_scale. However, it requires a specific software version and seems to have some inexplicable errors, which makes debugging difficult for many people who do not know CUDA programming.
If Pyg could support this kind of aggregation, it would be a great joy for the community! Thanks in advance.


### Additional context

_No response_


Discussion:
Yeah, I always thought this is super challenging to implement in an efficient way. Thanks for the pointers! We can try to look into it but it currently has lower priority.
Maybe I'm wrong, but I don't think we can break the $\log(n)$ depth complexity of parallel (arg)sort (that is basically the time complexity in a parallel setting). The (sequential) time complexity of the [selection algorithm](https://en.wikipedia.org/wiki/Selection_algorithm) for computing the median is $O(n)$, and even doing an approximation using median of medians we would obtain a log-depth complexity. So, I think we could use sorting with no hesitation (btw, I don't think the slow part in the [official `MedianGCN` implementation](https://github.com/EdisonLeeeee/MedianGCN) is really the argsort, but rather the fact that they compute a masking and a median [for every different node degree](https://github.com/EdisonLeeeee/MedianGCN/blob/b1f3db62bf378935624f3f23cc31e3ce83f8908d/median_pyg.py#L94), which could become very time expensive).

I tried implementing the `MedianAggregation` using indeed stable sorting, which is available in the newest versions of PyTorch (I used a similar trick in an [implementation](https://gist.github.com/flandolfi/6df01e8b04b4093c2dddf2fd3629ff1d) of the top-$k$ aggregation). Here is the code:
```python
from typing import Optional

import torch
from torch import Tensor

from torch_geometric.nn.aggr.base import Aggregation


class MedianAggregation(Aggregation):
    def forward(self, x: Tensor, index: Optional[Tensor] = None,
                ptr: Optional[Tensor] = None, dim_size: Optional[int] = None,
                dim: int = -2) -> Tensor:
        self.assert_index_present(index)  # we could also expand `ptr` to an index
        
        if dim_size is None:
            dim_size = index.max() + 1
        
        # compute the two median indices (they coincide if the number of elements is odd)
        count = torch.bincount(index, minlength=dim_size)
        cumsum = torch.cumsum(count, dim=0)
        l_idx = cumsum - count.div(2, rounding_mode='trunc') - 1
        r_idx = cumsum - (count + 1).div(2, rounding_mode='trunc')

        # shape used for expansions 
        # (1, ..., 1, -1, 1, ..., 1)
        shape = [1] * x.dim()
        shape[dim] = -1

        # two sorts: the first one on the value, the second (stable) on the indices
        x, x_perm = torch.sort(x, dim=dim)
        index = index.view(*shape).expand_as(x)
        index = index.take_along_dim(x_perm, dim=dim)
        index, index_perm = torch.sort(index, dim=dim, stable=True)
        x = x.take_along_dim(index_perm, dim=dim)

        # median as an average of the two central elements
        l_med = x.index_select(dim, l_idx)
        r_med = x.index_select(dim, r_idx)

        # if the number of elements is 0, return 'nan'
        out_mask = (count > 0).view(*shape)
        
        return torch.where(out_mask, 0.5*l_med + 0.5*r_med, float('nan'))
```
I did not made a proper testing, only some examples on the terminal:
```python
In [10]: x = torch.randint(10, size=(10, 5))
    ...: idx, _  = torch.randint(4, size=(10,)).sort()
    ...: aggr = MedianAggregation()
    ...: x, idx, aggr(x, idx, dim=0)
Out[10]: 
(tensor([[4, 6, 6, 3, 5],
         [3, 8, 1, 8, 3],
         [0, 4, 1, 6, 8],
         [0, 1, 9, 2, 1],
         [7, 3, 4, 9, 5],
         [4, 2, 1, 0, 9],
         [3, 6, 4, 3, 4],
         [7, 3, 7, 2, 9],
         [4, 9, 8, 7, 4],
         [5, 9, 3, 4, 4]]),
 tensor([0, 0, 0, 1, 1, 2, 2, 2, 3, 3]),
 tensor([[3.0000, 6.0000, 1.0000, 6.0000, 5.0000],
         [3.5000, 2.0000, 6.5000, 5.5000, 3.0000],
         [4.0000, 3.0000, 4.0000, 2.0000, 9.0000],
         [4.5000, 9.0000, 5.5000, 5.5000, 4.0000]]))
```
I think this is pretty cool! I haven't verified the correctness yet, but happy to do that if you are interested in contributing this implementation. Let me know!
Sure, I'll willingly self-assign. But I have a question: should I make a PR in PyG or in torch-scatter (as a "composite", maybe)? In the second case I'll have to wait to be public in pip/conda before putting an aggregation in PyG. Let me know!
Let's go with the first option and integrate it directly into PyG :)
Perfect, I'll do a PR as soon as possible!

Resolving patch:
diff --git a/torch_geometric/nn/aggr/__init__.py b/torch_geometric/nn/aggr/__init__.py
--- a/torch_geometric/nn/aggr/__init__.py
+++ b/torch_geometric/nn/aggr/__init__.py
@@ -11,6 +11,7 @@
     SoftmaxAggregation,
     PowerMeanAggregation,
 )
+from .quantile import MedianAggregation, QuantileAggregation
 from .lstm import LSTMAggregation
 from .set2set import Set2Set
 from .scaler import DegreeScalerAggregation
@@ -31,6 +32,8 @@
     'StdAggregation',
     'SoftmaxAggregation',
     'PowerMeanAggregation',
+    'MedianAggregation',
+    'QuantileAggregation',
     'LSTMAggregation',
     'Set2Set',
     'DegreeScalerAggregation',
diff --git a/torch_geometric/nn/aggr/quantile.py b/torch_geometric/nn/aggr/quantile.py
new file mode 100644
--- /dev/null
+++ b/torch_geometric/nn/aggr/quantile.py
@@ -0,0 +1,151 @@
+from typing import List, Optional, Union
+
+import torch
+from torch import Tensor
+
+from torch_geometric.nn.aggr import Aggregation
+
+
+class QuantileAggregation(Aggregation):
+    r"""An aggregation operator that returns the feature-wise :math:`q`-th
+    quantile of a set :math:`\mathcal{X}`. That is, for every feature
+    :math:`d`, it computes
+
+    .. math::
+        {\mathrm{Q}_q(\mathcal{X})}_d = \begin{cases}
+            x_{\pi_i,d} & i = q \cdot n, \\
+            f(x_{\pi_i,d}, x_{\pi_{i+1},d}) & i < q \cdot n < i + 1,\\
+        \end{cases}
+
+    where :math:`x_{\pi_1,d} \le \dots \le x_{\pi_i,d} \le \dots \le
+    x_{\pi_n,d}` and :math:`f(a, b)` is an interpolation
+    function defined by :obj:`interpolation`.
+
+    Args:
+        q (float or list): The quantile value(s) :math:`q`. Can be a scalar or
+            a list of scalars in the range :math:`[0, 1]`. If more than a
+            quantile is passed, the results are concatenated.
+        interpolation (str): Interpolation method applied if the quantile point
+            :math:`q\cdot n` lies between two values
+            :math:`a \le b`. Can be one of the following:
+
+            * :obj:`"lower"`: Returns the one with lowest value.
+
+            * :obj:`"higher"`: Returns the one with highest value.
+
+            * :obj:`"midpoint"`: Returns the average of the two values.
+
+            * :obj:`"nearest"`: Returns the one whose index is nearest to the
+              quantile point.
+
+            * :obj:`"linear"`: Returns a linear combination of the two
+              elements, defined as
+              :math:`f(a, b) = a + (b - a)\cdot(q\cdot n - i)`.
+
+            (default: :obj:`"linear"`)
+        fill_value (float, optional): The default value in the case no entry is
+            found for a given index (default: :obj:`0.0`).
+    """
+    interpolations = {'linear', 'lower', 'higher', 'nearest', 'midpoint'}
+
+    def __init__(self, q: Union[float, List[float]],
+                 interpolation: str = 'linear', fill_value: float = 0.0):
+        super().__init__()
+
+        qs = [q] if not isinstance(q, (list, tuple)) else q
+        if len(qs) == 0:
+            raise ValueError("Provide at least one quantile value for `q`.")
+        if not all(0. <= quantile <= 1. for quantile in qs):
+            raise ValueError("`q` must be in the range [0, 1].")
+        if interpolation not in self.interpolations:
+            raise ValueError(f"Invalid interpolation method "
+                             f"got ('{interpolation}')")
+
+        self._q = q
+        self.register_buffer('q', torch.Tensor(qs).view(-1, 1))
+        self.interpolation = interpolation
+        self.fill_value = fill_value
+
+    def forward(self, x: Tensor, index: Optional[Tensor] = None,
+                ptr: Optional[Tensor] = None, dim_size: Optional[int] = None,
+                dim: int = -2) -> Tensor:
+
+        dim = x.dim() + dim if dim < 0 else dim
+
+        self.assert_index_present(index)
+        assert index is not None  # Required for TorchScript.
+
+        count = torch.bincount(index, minlength=dim_size or 0)
+        cumsum = torch.cumsum(count, dim=0) - count
+
+        q_point = self.q * (count - 1) + cumsum
+        q_point = q_point.t().reshape(-1)
+
+        shape = [1] * x.dim()
+        shape[dim] = -1
+        index = index.view(shape).expand_as(x)
+
+        # Two sorts: the first one on the value,
+        # the second (stable) on the indices:
+        x, x_perm = torch.sort(x, dim=dim)
+        index = index.take_along_dim(x_perm, dim=dim)
+        index, index_perm = torch.sort(index, dim=dim, stable=True)
+        x = x.take_along_dim(index_perm, dim=dim)
+
+        # Compute the quantile interpolations:
+        if self.interpolation == 'lower':
+            quantile = x.index_select(dim, q_point.floor().long())
+        elif self.interpolation == 'higher':
+            quantile = x.index_select(dim, q_point.ceil().long())
+        elif self.interpolation == 'nearest':
+            quantile = x.index_select(dim, q_point.round().long())
+        else:
+            l_quant = x.index_select(dim, q_point.floor().long())
+            r_quant = x.index_select(dim, q_point.ceil().long())
+
+            if self.interpolation == 'linear':
+                q_frac = q_point.frac().view(shape)
+                quantile = l_quant + (r_quant - l_quant) * q_frac
+            else:  # 'midpoint'
+                quantile = 0.5 * l_quant + 0.5 * r_quant
+
+        # If the number of elements is zero, fill with pre-defined value:
+        mask = (count == 0).repeat_interleave(self.q.numel()).view(shape)
+        out = quantile.masked_fill(mask, self.fill_value)
+
+        if self.q.numel() > 1:
+            shape = list(out.shape)
+            shape = (shape[:dim] + [shape[dim] // self.q.numel(), -1] +
+                     shape[dim + 2:])
+            out = out.view(shape)
+
+        return out
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}(q={self._q})')
+
+
+class MedianAggregation(QuantileAggregation):
+    r"""An aggregation operator that returns the feature-wise median of a set.
+    That is, for every feature :math:`d`, it computes
+
+    .. math::
+        {\mathrm{median}(\mathcal{X})}_d = x_{\pi_i,d}
+
+    where :math:`x_{\pi_1,d} \le x_{\pi_2,d} \le \dots \le
+    x_{\pi_n,d}` and :math:`i = \lfloor \frac{n}{2} \rfloor`.
+
+    .. note::
+        If the median lies between two values, the lowest one is returned.
+        To compute the midpoint (or other kind of interpolation) of the two
+        values, use :class:`QuantileAggregation` instead.
+
+    Args:
+        fill_value (float, optional): The default value in the case no entry is
+            found for a given index (default: :obj:`0.0`).
+    """
+    def __init__(self, fill_value: float = 0.0):
+        super().__init__(0.5, 'lower', fill_value)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}()"


Test patch:
diff --git a/test/nn/aggr/test_quantile.py b/test/nn/aggr/test_quantile.py
new file mode 100644
--- /dev/null
+++ b/test/nn/aggr/test_quantile.py
@@ -0,0 +1,103 @@
+import pytest
+import torch
+
+from torch_geometric.nn import MedianAggregation, QuantileAggregation
+
+
+@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):
+    x = torch.tensor([
+        [0.0, 1.0, 2.0],
+        [3.0, 4.0, 5.0],
+        [6.0, 7.0, 8.0],
+        [9.0, 0.0, 1.0],
+        [2.0, 3.0, 4.0],
+        [5.0, 6.0, 7.0],
+        [8.0, 9.0, 0.0],
+        [1.0, 2.0, 3.0],
+        [4.0, 5.0, 6.0],
+        [7.0, 8.0, 9.0],
+    ])
+    index = torch.zeros(x.size(dim), dtype=torch.long)
+
+    aggr = QuantileAggregation(q=q, interpolation=interpolation)
+    assert str(aggr) == f"QuantileAggregation(q={q})"
+
+    out = aggr(x, index, dim=dim)
+    expected = x.quantile(q, dim, interpolation=interpolation, keepdim=True)
+    assert torch.allclose(out, expected)
+
+
+def test_median_aggregation():
+    x = torch.tensor([
+        [0.0, 1.0, 2.0],
+        [3.0, 4.0, 5.0],
+        [6.0, 7.0, 8.0],
+        [9.0, 0.0, 1.0],
+        [2.0, 3.0, 4.0],
+        [5.0, 6.0, 7.0],
+        [8.0, 9.0, 0.0],
+        [1.0, 2.0, 3.0],
+        [4.0, 5.0, 6.0],
+        [7.0, 8.0, 9.0],
+    ])
+
+    aggr = MedianAggregation()
+    assert str(aggr) == "MedianAggregation()"
+
+    index = torch.tensor([0, 0, 0, 0, 1, 1, 1, 2, 2, 2])
+    assert aggr(x, index).tolist() == [
+        [3.0, 1.0, 2.0],
+        [5.0, 6.0, 4.0],
+        [4.0, 5.0, 6.0],
+    ]
+
+    index = torch.tensor([0, 1, 0])
+    assert aggr(x, index, dim=1).tolist() == [
+        [0.0, 1.0],
+        [3.0, 4.0],
+        [6.0, 7.0],
+        [1.0, 0.0],
+        [2.0, 3.0],
+        [5.0, 6.0],
+        [0.0, 9.0],
+        [1.0, 2.0],
+        [4.0, 5.0],
+        [7.0, 8.0],
+    ]
+
+
+def test_quantile_aggregation_multi():
+    x = torch.tensor([
+        [0.0, 1.0, 2.0],
+        [3.0, 4.0, 5.0],
+        [6.0, 7.0, 8.0],
+        [9.0, 0.0, 1.0],
+        [2.0, 3.0, 4.0],
+        [5.0, 6.0, 7.0],
+        [8.0, 9.0, 0.0],
+        [1.0, 2.0, 3.0],
+        [4.0, 5.0, 6.0],
+        [7.0, 8.0, 9.0],
+    ])
+    index = torch.tensor([0, 0, 0, 0, 1, 1, 1, 2, 2, 2])
+
+    qs = [0.25, 0.5, 0.75]
+
+    assert torch.allclose(
+        QuantileAggregation(qs)(x, index),
+        torch.cat([QuantileAggregation(q)(x, index) for q in qs], dim=-1),
+    )
+
+
+def test_quantile_aggregation_validate():
+    with pytest.raises(ValueError, match="at least one quantile"):
+        QuantileAggregation(q=[])
+
+    with pytest.raises(ValueError, match="must be in the range"):
+        QuantileAggregation(q=-1)
+
+    with pytest.raises(ValueError, match="Invalid interpolation method"):
+        QuantileAggregation(q=0.5, interpolation=None)