Repository: pyg-team/pytorch_geometric
Pull request #7339 (2023-05-10 10:39:37)

Add `add_pad_mask` argument to `Pad` transform
add_masks_to_data is an optional argument that will cause the addition of two masks to data object. They will indicate which elements in the data are real and which are added as a padding.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7339?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7339](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7339?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (4cf6a03) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/b2ee130059e7faff065dc9b8df632595e0724e18?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (b2ee130) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #7339   +/-   ##
=======================================
  Coverage   91.52%   91.52%           
=======================================
  Files         438      438           
  Lines       24346    24353    +7     
=======================================
+ Hits        22283    22290    +7     
  Misses       2063     2063           
```


| [Impacted Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7339?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/transforms/pad.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7339?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3RyYW5zZm9ybXMvcGFkLnB5) | `95.73% <100.00%> (+0.14%)` | :arrow_up: |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3f97b8446..bea641c9a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added an optional `add_pad_mask` argument to the `Pad` transform ([#7339](https://github.com/pyg-team/pytorch_geometric/pull/7339))
 - Added `keep_inter_cluster_edges` option to `ClusterData` to support inter-subgraph edge connections when doing graph partitioning ([#7326](https://github.com/pyg-team/pytorch_geometric/pull/7326))
 - Unify graph pooling framework ([#7308](https://github.com/pyg-team/pytorch_geometric/pull/7308))
 - Added support for tuples as keys in `ModuleDict`/`ParameterDict` ([#7294](https://github.com/pyg-team/pytorch_geometric/pull/7294))
diff --git a/torch_geometric/transforms/pad.py b/torch_geometric/transforms/pad.py
index f9fff7a66..b92b8a983 100644
--- a/torch_geometric/transforms/pad.py
+++ b/torch_geometric/transforms/pad.py
@@ -265,6 +265,12 @@ class Pad(BaseTransform):
         mask_pad_value (bool, optional): The fill value to use for
             :obj:`train_mask`, :obj:`val_mask` and :obj:`test_mask` attributes
             (default: :obj:`False`).
+        add_pad_mask (bool, optional): If set to :obj:`True`, will attach
+            node-level :obj:`pad_node_mask` and edge-level :obj:`pad_edge_mask`
+            attributes to the output which indicates which elements in the data
+            are real (represented by :obj:`True`) and which were added as a
+            result of padding (represented by :obj:`False`).
+            (default: :obj:`False`)
         exclude_keys ([str], optional): Keys to be removed
             from the input data object. (default: :obj:`None`)
     """
@@ -275,6 +281,7 @@ class Pad(BaseTransform):
         node_pad_value: Union[int, float, Padding] = 0.0,
         edge_pad_value: Union[int, float, Padding] = 0.0,
         mask_pad_value: bool = False,
+        add_pad_mask: bool = False,
         exclude_keys: Optional[List[str]] = None,
     ):
         self.max_num_nodes = self._NumNodes(max_num_nodes)
@@ -293,6 +300,7 @@ class Pad(BaseTransform):
             for key in ['train_mask', 'val_mask', 'test_mask']
         }
 
+        self.add_pad_mask = add_pad_mask
         self.exclude_keys = set(exclude_keys or [])
 
     class _IntOrDict(ABC):
@@ -448,6 +456,11 @@ class Pad(BaseTransform):
             f'({store.num_nodes}).'
         num_pad_nodes = num_target_nodes - store.num_nodes
 
+        if self.add_pad_mask:
+            pad_node_mask = torch.ones(num_target_nodes, dtype=torch.bool)
+            pad_node_mask[store.num_nodes:] = False
+            store.pad_node_mask = pad_node_mask
+
         for attr_name in attrs_to_pad:
             attr = store[attr_name]
             pad_value = self.__get_node_padding(attr_name, node_type)
@@ -470,6 +483,11 @@ class Pad(BaseTransform):
             f'({store.num_edges}).'
         num_pad_edges = num_target_edges - store.num_edges
 
+        if self.add_pad_mask:
+            pad_edge_mask = torch.ones(num_target_edges, dtype=torch.bool)
+            pad_edge_mask[store.num_edges:] = False
+            store.pad_edge_mask = pad_edge_mask
+
         if isinstance(num_nodes, tuple):
             src_pad_value, dst_pad_value = num_nodes
         else:


Test patch:
diff --git a/test/transforms/test_pad.py b/test/transforms/test_pad.py
index e1fea40ce..69225a942 100644
--- a/test/transforms/test_pad.py
+++ b/test/transforms/test_pad.py
@@ -56,12 +56,23 @@ def _generate_heterodata_edges(
             yield edge_type, attr
 
 
-def _check_homo_data_nodes(original: Data, padded: Data,
-                           max_num_nodes: Union[int, Dict[NodeType, int]],
-                           node_pad_value: Optional[Padding] = None,
-                           exclude_keys: Optional[List[str]] = None):
+def _check_homo_data_nodes(
+    original: Data,
+    padded: Data,
+    max_num_nodes: Union[int, Dict[NodeType, int]],
+    node_pad_value: Optional[Padding] = None,
+    is_mask_available: bool = False,
+    exclude_keys: Optional[List[str]] = None,
+):
     assert padded.num_nodes == max_num_nodes
 
+    compare_pad_start_idx = original.num_nodes
+
+    if is_mask_available:
+        assert padded.pad_node_mask.numel() == padded.num_nodes
+        assert torch.all(padded.pad_node_mask[:compare_pad_start_idx])
+        assert not torch.any(padded.pad_node_mask[compare_pad_start_idx:])
+
     for attr in _generate_homodata_node_attrs(original):
         if attr in exclude_keys:
             assert attr not in padded.keys
@@ -73,7 +84,7 @@ def _check_homo_data_nodes(original: Data, padded: Data,
             continue
 
         assert padded[attr].shape[0] == max_num_nodes
-        compare_pad_start_idx = original[attr].shape[0]
+
         # Check values in padded area.
         pad_value = node_pad_value.get_value(
             None, attr) if node_pad_value is not None else 0.0
@@ -86,10 +97,14 @@ def _check_homo_data_nodes(original: Data, padded: Data,
                            padded[attr][:compare_pad_start_idx])
 
 
-def _check_homo_data_edges(original: Data, padded: Data,
-                           max_num_edges: Optional[int] = None,
-                           edge_pad_value: Optional[Padding] = None,
-                           exclude_keys: Optional[List[str]] = None):
+def _check_homo_data_edges(
+    original: Data,
+    padded: Data,
+    max_num_edges: Optional[int] = None,
+    edge_pad_value: Optional[Padding] = None,
+    is_mask_available: bool = False,
+    exclude_keys: Optional[List[str]] = None,
+):
     # Check edge index attribute.
     if max_num_edges is None:
         max_num_edges = padded.num_nodes**2
@@ -108,6 +123,11 @@ def _check_homo_data_edges(original: Data, padded: Data,
     assert torch.equal(original.edge_index,
                        padded.edge_index[:, :compare_pad_start_idx])
 
+    if is_mask_available:
+        assert padded.pad_edge_mask.numel() == padded.num_edges
+        assert torch.all(padded.pad_edge_mask[:compare_pad_start_idx])
+        assert not torch.any(padded.pad_edge_mask[compare_pad_start_idx:])
+
     # Check other attributes.
     for attr in _generate_homodata_edge_attrs(original):
         if attr == 'edge_index':
@@ -135,10 +155,17 @@ def _check_homo_data_edges(original: Data, padded: Data,
                            padded[attr][:compare_pad_start_idx, :])
 
 
-def _check_hetero_data_nodes(original: HeteroData, padded: HeteroData,
-                             max_num_nodes: Union[int, Dict[NodeType, int]],
-                             node_pad_value: Optional[Padding] = None,
-                             exclude_keys: Optional[List[str]] = None):
+def _check_hetero_data_nodes(
+    original: HeteroData,
+    padded: HeteroData,
+    max_num_nodes: Union[int, Dict[NodeType, int]],
+    node_pad_value: Optional[Padding] = None,
+    is_mask_available: bool = False,
+    exclude_keys: Optional[List[str]] = None,
+):
+    if is_mask_available:
+        for store in padded.node_stores:
+            assert 'pad_node_mask' in store
 
     expected_nodes = max_num_nodes
 
@@ -152,9 +179,17 @@ def _check_hetero_data_nodes(original: HeteroData, padded: HeteroData,
         if not isinstance(padded[node_type][attr], torch.Tensor):
             continue
 
-        original_tensor = original[node_type][attr]
+        compare_pad_start_idx = original[node_type].num_nodes
         padded_tensor = padded[node_type][attr]
 
+        if attr == 'pad_node_mask':
+            assert padded_tensor.numel() == padded[node_type].num_nodes
+            assert torch.all(padded_tensor[:compare_pad_start_idx])
+            assert not torch.any(padded_tensor[compare_pad_start_idx:])
+            continue
+
+        original_tensor = original[node_type][attr]
+
         # Check the number of nodes.
         if isinstance(max_num_nodes, dict):
             expected_nodes = max_num_nodes[node_type]
@@ -172,12 +207,17 @@ def _check_hetero_data_nodes(original: HeteroData, padded: HeteroData,
                            padded_tensor[:compare_pad_start_idx])
 
 
-def _check_hetero_data_edges(original: HeteroData, padded: HeteroData,
-                             max_num_edges: Optional[Union[int,
-                                                           Dict[EdgeType,
-                                                                int]]] = None,
-                             edge_pad_value: Optional[Padding] = None,
-                             exclude_keys: Optional[List[str]] = None):
+def _check_hetero_data_edges(
+    original: HeteroData,
+    padded: HeteroData,
+    max_num_edges: Optional[Union[int, Dict[EdgeType, int]]] = None,
+    edge_pad_value: Optional[Padding] = None,
+    is_mask_available: bool = False,
+    exclude_keys: Optional[List[str]] = None,
+):
+    if is_mask_available:
+        for store in padded.edge_stores:
+            assert 'pad_edge_mask' in store
 
     for edge_type, attr in _generate_heterodata_edges(padded):
         if attr in exclude_keys:
@@ -190,9 +230,16 @@ def _check_hetero_data_edges(original: HeteroData, padded: HeteroData,
             continue
 
         compare_pad_start_idx = original[edge_type].num_edges
-        original_tensor = original[edge_type][attr]
         padded_tensor = padded[edge_type][attr]
 
+        if attr == 'pad_edge_mask':
+            assert padded_tensor.numel() == padded[edge_type].num_edges
+            assert torch.all(padded_tensor[:compare_pad_start_idx])
+            assert not torch.any(padded_tensor[compare_pad_start_idx:])
+            continue
+
+        original_tensor = original[edge_type][attr]
+
         if isinstance(max_num_edges, numbers.Number):
             expected_num_edges = max_num_edges
         elif max_num_edges is None or edge_type not in max_num_edges.keys():
@@ -235,33 +282,40 @@ def _check_hetero_data_edges(original: HeteroData, padded: HeteroData,
                                padded_tensor[:compare_pad_start_idx, :])
 
 
-def _check_data(original: Union[Data, HeteroData], padded: Union[Data,
-                                                                 HeteroData],
-                max_num_nodes: Union[int, Dict[NodeType, int]],
-                max_num_edges: Optional[Union[int, Dict[EdgeType,
-                                                        int]]] = None,
-                node_pad_value: Optional[Union[Padding, int, float]] = None,
-                edge_pad_value: Optional[Union[Padding, int, float]] = None,
-                exclude_keys: Optional[List[str]] = None):
+def _check_data(
+    original: Union[Data, HeteroData],
+    padded: Union[Data, HeteroData],
+    max_num_nodes: Union[int, Dict[NodeType, int]],
+    max_num_edges: Optional[Union[int, Dict[EdgeType, int]]] = None,
+    node_pad_value: Optional[Union[Padding, int, float]] = None,
+    edge_pad_value: Optional[Union[Padding, int, float]] = None,
+    is_mask_available: bool = False,
+    exclude_keys: Optional[List[str]] = None,
+):
 
     if not isinstance(node_pad_value, Padding) and node_pad_value is not None:
         node_pad_value = UniformPadding(node_pad_value)
     if not isinstance(edge_pad_value, Padding) and edge_pad_value is not None:
         edge_pad_value = UniformPadding(edge_pad_value)
 
+    if is_mask_available is None:
+        is_mask_available = False
+
     if exclude_keys is None:
         exclude_keys = []
 
     if isinstance(original, Data):
         _check_homo_data_nodes(original, padded, max_num_nodes, node_pad_value,
-                               exclude_keys)
+                               is_mask_available, exclude_keys)
         _check_homo_data_edges(original, padded, max_num_edges, edge_pad_value,
-                               exclude_keys)
+                               is_mask_available, exclude_keys)
     else:
         _check_hetero_data_nodes(original, padded, max_num_nodes,
-                                 node_pad_value, exclude_keys)
+                                 node_pad_value, is_mask_available,
+                                 exclude_keys)
         _check_hetero_data_edges(original, padded, max_num_edges,
-                                 edge_pad_value, exclude_keys)
+                                 edge_pad_value, is_mask_available,
+                                 exclude_keys)
 
 
 def test_pad_repr():
@@ -273,35 +327,42 @@ def test_pad_repr():
 
 @pytest.mark.parametrize('data', [fake_data(), fake_hetero_data()])
 @pytest.mark.parametrize('num_nodes', [32, 64])
-def test_pad_auto_edges(data, num_nodes):
+@pytest.mark.parametrize('add_pad_mask', [True, False])
+def test_pad_auto_edges(data, num_nodes, add_pad_mask):
     original = data
     data = deepcopy(data)
-    transform = Pad(max_num_nodes=num_nodes)
+    transform = Pad(max_num_nodes=num_nodes, add_pad_mask=add_pad_mask)
 
     padded = transform(data)
-    _check_data(original, padded, num_nodes)
+    _check_data(original, padded, num_nodes, is_mask_available=add_pad_mask)
 
 
 @pytest.mark.parametrize('num_nodes', [32, 64])
 @pytest.mark.parametrize('num_edges', [300, 411])
-def test_pad_data_explicit_edges(num_nodes, num_edges):
+@pytest.mark.parametrize('add_pad_mask', [True, False])
+def test_pad_data_explicit_edges(num_nodes, num_edges, add_pad_mask):
     data = fake_data()
     original = deepcopy(data)
-    transform = Pad(max_num_nodes=num_nodes, max_num_edges=num_edges)
+    transform = Pad(max_num_nodes=num_nodes, max_num_edges=num_edges,
+                    add_pad_mask=add_pad_mask)
 
     padded = transform(data)
-    _check_data(original, padded, num_nodes, num_edges)
+    _check_data(original, padded, num_nodes, num_edges,
+                is_mask_available=add_pad_mask)
 
 
 @pytest.mark.parametrize('num_nodes', [32, {'v0': 64, 'v1': 36}])
 @pytest.mark.parametrize('num_edges', [300, {('v0', 'e0', 'v1'): 397}])
-def test_pad_heterodata_explicit_edges(num_nodes, num_edges):
+@pytest.mark.parametrize('add_pad_mask', [True, False])
+def test_pad_heterodata_explicit_edges(num_nodes, num_edges, add_pad_mask):
     data = fake_hetero_data()
     original = deepcopy(data)
-    transform = Pad(max_num_nodes=num_nodes, max_num_edges=num_edges)
+    transform = Pad(max_num_nodes=num_nodes, max_num_edges=num_edges,
+                    add_pad_mask=add_pad_mask)
 
     padded = transform(data)
-    _check_data(original, padded, num_nodes, num_edges)
+    _check_data(original, padded, num_nodes, num_edges,
+                is_mask_available=add_pad_mask)
 
 
 @pytest.mark.parametrize('node_pad_value', [10, AttrNamePadding({'x': 3.0})])
@@ -348,16 +409,22 @@ def test_pad_heterodata_pad_values(node_pad_value, edge_pad_value):
 
 
 @pytest.mark.parametrize('data', [fake_data(), fake_hetero_data()])
-@pytest.mark.parametrize('exclude_keys',
-                         [['y'], ['edge_attr'], ['y', 'edge_attr']])
-def test_pad_data_exclude_keys(data, exclude_keys):
+@pytest.mark.parametrize('add_pad_mask', [True, False])
+@pytest.mark.parametrize('exclude_keys', [
+    ['y'],
+    ['edge_attr'],
+    ['y', 'edge_attr'],
+])
+def test_pad_data_exclude_keys(data, add_pad_mask, exclude_keys):
     original = data
     data = deepcopy(data)
     num_nodes = 32
-    transform = Pad(max_num_nodes=num_nodes, exclude_keys=exclude_keys)
+    transform = Pad(max_num_nodes=num_nodes, add_pad_mask=add_pad_mask,
+                    exclude_keys=exclude_keys)
 
     padded = transform(data)
-    _check_data(original, padded, num_nodes, exclude_keys=exclude_keys)
+    _check_data(original, padded, num_nodes, is_mask_available=add_pad_mask,
+                exclude_keys=exclude_keys)
 
 
 @pytest.mark.parametrize('data', [fake_data(), fake_hetero_data(node_types=1)])