Repository: pyg-team/pytorch_geometric

Issue:
Aggregation.__call__ calls nn.Module.__call__ using (x, index, ptr, dim_size, dim, **kwargs) instead of (x=x, index=index, ptr=ptr, dim_size=dim_size, dim=dim, **kwargs) 
### 🛠 Proposed Refactor

if somebody wants to inheric Aggregation base class, AND wants to add some conditional data used in aggregation, he must re-define forward as (x: Tensor, index: Tensor=None, ptr=None, dim_size: int=None, dim: int = -2, **condition: Tensor=None**), since the internal call of Aggregation to nn.Module uses fixed positional args. I think it is more elegant to define forward as (x, condition, index=None, ptr=None, dim_size=None, dim=None)

### Suggest a potential alternative/fix

As suggested


Discussion:

Resolving patch:
diff --git a/torch_geometric/nn/aggr/base.py b/torch_geometric/nn/aggr/base.py
--- a/torch_geometric/nn/aggr/base.py
+++ b/torch_geometric/nn/aggr/base.py
@@ -93,9 +93,15 @@ def reset_parameters(self):
         pass
 
     @disable_dynamic_shapes(required_args=['dim_size'])
-    def __call__(self, x: Tensor, index: Optional[Tensor] = None,
-                 ptr: Optional[Tensor] = None, dim_size: Optional[int] = None,
-                 dim: int = -2, **kwargs) -> Tensor:
+    def __call__(
+        self,
+        x: Tensor,
+        index: Optional[Tensor] = None,
+        ptr: Optional[Tensor] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        **kwargs,
+    ) -> Tensor:
 
         if dim >= x.dim() or dim < -x.dim():
             raise ValueError(f"Encountered invalid dimension '{dim}' of "
@@ -116,7 +122,8 @@ def __call__(self, x: Tensor, index: Optional[Tensor] = None,
             dim_size = int(index.max()) + 1 if index.numel() > 0 else 0
 
         try:
-            return super().__call__(x, index, ptr, dim_size, dim, **kwargs)
+            return super().__call__(x, index=index, ptr=ptr, dim_size=dim_size,
+                                    dim=dim, **kwargs)
         except (IndexError, RuntimeError) as e:
             if index is not None:
                 if index.numel() > 0 and dim_size <= int(index.max()):
diff --git a/torch_geometric/sampler/neighbor_sampler.py b/torch_geometric/sampler/neighbor_sampler.py
--- a/torch_geometric/sampler/neighbor_sampler.py
+++ b/torch_geometric/sampler/neighbor_sampler.py
@@ -217,13 +217,12 @@ def _sample(
             # TODO Support induced subgraph sampling in `pyg-lib`.
             if (torch_geometric.typing.WITH_PYG_LIB
                     and self.subgraph_type != SubgraphType.induced):
-                # TODO (matthias) `return_edge_id` if edge features present
                 # TODO (matthias) Ideally, `seed` inherits dtype from `colptr`
                 colptrs = list(self.colptr_dict.values())
                 dtype = colptrs[0].dtype if len(colptrs) > 0 else torch.int64
                 seed = {k: v.to(dtype) for k, v in seed.items()}
 
-                out = torch.ops.pyg.hetero_neighbor_sample(
+                args = (
                     self.node_types,
                     self.edge_types,
                     self.colptr_dict,
@@ -232,13 +231,20 @@ def _sample(
                     self.num_neighbors.get_mapped_values(self.edge_types),
                     self.node_time,
                     seed_time,
+                )
+                if torch_geometric.typing.WITH_WEIGHTED_NEIGHBOR_SAMPLE:
+                    args += (None, )
+                args += (
                     True,  # csc
                     self.replace,
                     self.subgraph_type != SubgraphType.induced,
                     self.disjoint,
                     self.temporal_strategy,
+                    # TODO (matthias) `return_edge_id` if edge features present
                     True,  # return_edge_id
                 )
+
+                out = torch.ops.pyg.hetero_neighbor_sample(*args)
                 row, col, node, edge, batch = out[:4] + (None, )
 
                 # `pyg-lib>0.1.0` returns sampled number of nodes/edges:
diff --git a/torch_geometric/transforms/random_link_split.py b/torch_geometric/transforms/random_link_split.py
--- a/torch_geometric/transforms/random_link_split.py
+++ b/torch_geometric/transforms/random_link_split.py
@@ -41,10 +41,10 @@ class RandomLinkSplit(BaseTransform):
             (default: :obj:`0.2`)
         is_undirected (bool): If set to :obj:`True`, the graph is assumed to be
             undirected, and positive and negative samples will not leak
-            (reverse) edge connectivity across different splits. Note that this
-            only affects the graph split, label data will not be returned
-            undirected.
-            (default: :obj:`False`)
+            (reverse) edge connectivity across different splits. This only
+            affects the graph split, label data will not be returned
+            undirected. This option is ignored for bipartite edge types or
+            whenever :obj:`edge_type != rev_edge_type`. (default: :obj:`False`)
         key (str, optional): The name of the attribute holding
             ground-truth labels.
             If :obj:`data[key]` does not exist, it will be automatically


Test patch:
diff --git a/test/loader/test_neighbor_loader.py b/test/loader/test_neighbor_loader.py
--- a/test/loader/test_neighbor_loader.py
+++ b/test/loader/test_neighbor_loader.py
@@ -564,8 +564,8 @@ def test_pyg_lib_and_torch_sparse_hetero_equality():
 
     sample = torch.ops.pyg.hetero_neighbor_sample
     out1 = sample(node_types, edge_types, colptr_dict, row_dict, seed_dict,
-                  num_neighbors_dict, None, None, True, False, True, False,
-                  "uniform", True)
+                  num_neighbors_dict, None, None, None, True, False, True,
+                  False, "uniform", True)
     sample = torch.ops.torch_sparse.hetero_neighbor_sample
     out2 = sample(node_types, edge_types, colptr_dict, row_dict, seed_dict,
                   num_neighbors_dict, 2, False, True)
diff --git a/test/nn/test_model_summary.py b/test/nn/test_model_summary.py
--- a/test/nn/test_model_summary.py
+++ b/test/nn/test_model_summary.py
@@ -148,10 +148,10 @@ def test_summary_with_leaf_module(gcn):
 | ├─(act)ReLU                             | [100, 16]          | [100, 16]      | --       |
 | ├─(convs)ModuleList                     | --                 | --             | 1,072    |
 | │    └─(0)GCNConv                       | [100, 32], [2, 20] | [100, 16]      | 528      |
-| │    │    └─(aggr_module)SumAggregation | [120, 16], [120]   | [100, 16]      | --       |
+| │    │    └─(aggr_module)SumAggregation | [120, 16]          | [100, 16]      | --       |
 | │    │    └─(lin)Linear                 | [100, 32]          | [100, 16]      | 512      |
 | │    └─(1)GCNConv                       | [100, 16], [2, 20] | [100, 32]      | 544      |
-| │    │    └─(aggr_module)SumAggregation | [120, 32], [120]   | [100, 32]      | --       |
+| │    │    └─(aggr_module)SumAggregation | [120, 32]          | [100, 32]      | --       |
 | │    │    └─(lin)Linear                 | [100, 16]          | [100, 32]      | 512      |
 | ├─(norms)ModuleList                     | --                 | --             | --       |
 | │    └─(0)Identity                      | [100, 16]          | [100, 16]      | --       |