Add init arguments in AddSelfLoops

test/transforms/test_add_self_loops.py

@@ -7,6 +7,8 @@ def test_add_self_loops():
     assert AddSelfLoops().__repr__() == 'AddSelfLoops()'
 
     edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
+    edge_weight = torch.tensor([1, 2, 3, 4])
+    edge_attr = torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8]])
 
     data = Data(edge_index=edge_index, num_nodes=3)
     data = AddSelfLoops()(data)
@@ -15,6 +17,21 @@ def test_add_self_loops():
                                         [1, 0, 2, 1, 0, 1, 2]]
     assert data.num_nodes == 3
 
+    data = Data(edge_index=edge_index, edge_weight=edge_weight, num_nodes=3)
+    data = AddSelfLoops(attr='edge_weight', fill_value=5)(data)
+    assert data.edge_index.tolist() == [[0, 1, 1, 2, 0, 1, 2],
+                                        [1, 0, 2, 1, 0, 1, 2]]
+    assert data.num_nodes == 3
+    assert data.edge_weight.tolist() == [1, 2, 3, 4, 5, 5, 5]
+
+    data = Data(edge_index=edge_index, edge_attr=edge_attr, num_nodes=3)
+    data = AddSelfLoops(attr='edge_attr', fill_value='add')(data)
+    assert data.edge_index.tolist() == [[0, 1, 1, 2, 0, 1, 2],
+                                        [1, 0, 2, 1, 0, 1, 2]]
+    assert data.num_nodes == 3
+    assert data.edge_attr.tolist() == [[1, 2], [3, 4], [5, 6], [7, 8],
+                                       [3, 4], [8, 10], [5, 6]]
+
 
 def test_hetero_add_self_loops():
     edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])

torch_geometric/transforms/add_self_loops.py

@@ -1,22 +1,47 @@
-from typing import Union
+from torch import Tensor
+
+from typing import Union, Optional
 
 from torch_geometric.utils import add_self_loops
 from torch_geometric.data import Data, HeteroData
 from torch_geometric.transforms import BaseTransform
 
 
 class AddSelfLoops(BaseTransform):
-    r"""Adds self-loops to the given homogeneous or heterogeneous graph."""
+    r"""Adds self-loops to the given homogeneous or heterogeneous graph.
+
+    Args:
+        attr: (str, optional): The name of the attribute of edge weights
+            or multi-dimensional edge features to pass to
+            :meth:`torch_geometric.utils.add_self_loops`.
+            (default: :obj:`"edge_weight"`)
+        fill_value (float or Tensor or str, optional): The way to generate
+            edge features of self-loops (in case :obj:`attr != None`).
+            If given as :obj:`float` or :class:`torch.Tensor`, edge features of
+            self-loops will be directly given by :obj:`fill_value`.
+            If given as :obj:`str`, edge features of self-loops are computed by
+            aggregating all features of edges that point to the specific node,
+            according to a reduce operation. (:obj:`"add"`, :obj:`"mean"`,
+            :obj:`"min"`, :obj:`"max"`, :obj:`"mul"`). (default: :obj:`1.`)
+    """
+    def __init__(self, attr: Optional[str] = 'edge_weight',
+                 fill_value: Union[float, Tensor, str] = None):
+        self.attr = attr
+        self.fill_value = fill_value
+
     def __call__(self, data: Union[Data, HeteroData]):
         for store in data.edge_stores:
             if store.is_bipartite() or 'edge_index' not in store:
                 continue
 
-            store.edge_index, store.edge_weight = add_self_loops(
+            store.edge_index, edge_weight = add_self_loops(
                 store.edge_index,
-                store.edge_weight if 'edge_weight' in store else None,
+                getattr(store, self.attr, None),
+                fill_value=self.fill_value,
                 num_nodes=store.size(0))
 
+            setattr(store, self.attr, edge_weight)
+
         return data
 
     def __repr__(self) -> str: