Repository: pyg-team/pytorch_geometric
Pull request #8622 (2023-12-14 20:25:58)

Support `mypy` in `torch_geometric.transforms` [6/6]

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8622?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`cd97a6e`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/cd97a6ee2102be143eed98dc0495c56ec45829d0?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.19% compared to head [(`5fc8ac7`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8622?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.19%.


<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##           master    #8622   +/-   ##
=======================================
  Coverage   89.19%   89.19%           
=======================================
  Files         481      481           
  Lines       30707    30708    +1     
=======================================
+ Hits        27388    27389    +1     
  Misses       3319     3319           
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8622?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).


Patch:
diff --git a/pyproject.toml b/pyproject.toml
index 94c26546b..08e65339f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -124,7 +124,6 @@ module = [
     "torch_geometric.data.*",
     "torch_geometric.sampler.*",
     "torch_geometric.loader.*",
-    "torch_geometric.transforms.*",
     "torch_geometric.nn.*",
     "torch_geometric.explain.*",
     "torch_geometric.profile.*",
diff --git a/torch_geometric/data/data.py b/torch_geometric/data/data.py
index 5db0021f4..8c8c6e528 100644
--- a/torch_geometric/data/data.py
+++ b/torch_geometric/data/data.py
@@ -19,6 +19,7 @@ from typing import (
 import numpy as np
 import torch
 from torch import Tensor
+from typing_extensions import Self
 
 from torch_geometric.data import EdgeAttr, FeatureStore, GraphStore, TensorAttr
 from torch_geometric.data.feature_store import _FieldStatus
@@ -70,7 +71,7 @@ class BaseData:
     def __repr__(self) -> str:
         raise NotImplementedError
 
-    def stores_as(self, data: 'BaseData'):
+    def stores_as(self, data: Self):
         raise NotImplementedError
 
     @property
@@ -93,13 +94,13 @@ class BaseData:
         r"""Returns a :obj:`NamedTuple` of stored key/value pairs."""
         raise NotImplementedError
 
-    def update(self, data: 'BaseData') -> 'BaseData':
+    def update(self, data: Self) -> Self:
         r"""Updates the data object with the elements from another data object.
         Added elements will override existing ones (in case of duplicates).
         """
         raise NotImplementedError
 
-    def concat(self, data: 'BaseData') -> 'BaseData':
+    def concat(self, data: Self) -> Self:
         r"""Concatenates :obj:`self` with another :obj:`data` object.
         All values needs to have matching shapes at non-concat dimensions.
         """
@@ -245,7 +246,7 @@ class BaseData:
         return all(
             [store.is_sorted(sort_by_row) for store in self.edge_stores])
 
-    def sort(self, sort_by_row: bool = True) -> 'Data':
+    def sort(self, sort_by_row: bool = True) -> Self:
         r"""Sorts edge indices :obj:`edge_index` and their corresponding edge
         features.
 
@@ -265,7 +266,7 @@ class BaseData:
         """
         return all([store.is_coalesced() for store in self.edge_stores])
 
-    def coalesce(self) -> 'Data':
+    def coalesce(self) -> Self:
         r"""Sorts and removes duplicated entries from edge indices
         :obj:`edge_index`.
         """
@@ -278,7 +279,7 @@ class BaseData:
         r"""Returns :obj:`True` if :obj:`time` is sorted."""
         return all([store.is_sorted_by_time() for store in self.stores])
 
-    def sort_by_time(self) -> 'BaseData':
+    def sort_by_time(self) -> Self:
         r"""Sorts data associated with :obj:`time` according to :obj:`time`."""
         out = copy.copy(self)
         for store in out.stores:
@@ -289,7 +290,7 @@ class BaseData:
         self,
         start_time: Union[float, int],
         end_time: Union[float, int],
-    ) -> 'BaseData':
+    ) -> Self:
         r"""Returns a snapshot of :obj:`data` to only hold events that occurred
         in period :obj:`[start_time, end_time]`.
         """
@@ -298,7 +299,7 @@ class BaseData:
             store.snapshot(start_time, end_time)
         return out
 
-    def up_to(self, end_time: Union[float, int]) -> 'BaseData':
+    def up_to(self, end_time: Union[float, int]) -> Self:
         r"""Returns a snapshot of :obj:`data` to only hold events that occurred
         up to :obj:`end_time` (inclusive of :obj:`edge_time`).
         """
@@ -616,7 +617,7 @@ class Data(BaseData, FeatureStore, GraphStore):
     def num_nodes(self, num_nodes: Optional[int]):
         self._store.num_nodes = num_nodes
 
-    def stores_as(self, data: 'Data'):
+    def stores_as(self, data: Self):
         return self
 
     @property
@@ -637,7 +638,7 @@ class Data(BaseData, FeatureStore, GraphStore):
     def to_namedtuple(self) -> NamedTuple:
         return self._store.to_namedtuple()
 
-    def update(self, data: Union['Data', Dict[str, Any]]) -> 'Data':
+    def update(self, data: Union[Self, Dict[str, Any]]) -> Self:
         for key, value in data.items():
             self[key] = value
         return self
@@ -709,7 +710,7 @@ class Data(BaseData, FeatureStore, GraphStore):
         """
         return self._store.is_edge_attr(key)
 
-    def subgraph(self, subset: Tensor) -> 'Data':
+    def subgraph(self, subset: Tensor) -> Self:
         r"""Returns the induced subgraph given by the node indices
         :obj:`subset`.
 
@@ -751,7 +752,7 @@ class Data(BaseData, FeatureStore, GraphStore):
 
         return data
 
-    def edge_subgraph(self, subset: Tensor) -> 'Data':
+    def edge_subgraph(self, subset: Tensor) -> Self:
         r"""Returns the induced subgraph given by the edge indices
         :obj:`subset`.
         Will currently preserve all the nodes in the graph, even if they are
@@ -896,7 +897,7 @@ class Data(BaseData, FeatureStore, GraphStore):
     ###########################################################################
 
     @classmethod
-    def from_dict(cls, mapping: Dict[str, Any]) -> 'Data':
+    def from_dict(cls, mapping: Dict[str, Any]) -> Self:
         r"""Creates a :class:`~torch_geometric.data.Data` object from a
         dictionary.
         """
diff --git a/torch_geometric/data/hetero_data.py b/torch_geometric/data/hetero_data.py
index bbffa3bba..53aa4369b 100644
--- a/torch_geometric/data/hetero_data.py
+++ b/torch_geometric/data/hetero_data.py
@@ -8,6 +8,7 @@ from typing import Any, Dict, List, NamedTuple, Optional, Tuple, Union
 
 import torch
 from torch import Tensor
+from typing_extensions import Self
 
 from torch_geometric.data import EdgeAttr, FeatureStore, GraphStore, TensorAttr
 from torch_geometric.data.data import BaseData, Data, size_repr, warn_or_raise
@@ -129,7 +130,7 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
                 setattr(self, key, value)
 
     @classmethod
-    def from_dict(cls, mapping: Dict[str, Any]) -> 'HeteroData':
+    def from_dict(cls, mapping: Dict[str, Any]) -> Self:
         r"""Creates a :class:`~torch_geometric.data.HeteroData` object from a
         dictionary.
         """
@@ -237,7 +238,7 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
         info = f'\n{info}\n' if len(info) > 0 else info
         return f'{self.__class__.__name__}({info})'
 
-    def stores_as(self, data: 'HeteroData'):
+    def stores_as(self, data: Self):
         for node_type in data.node_types:
             self.get_node_store(node_type)
         for edge_type in data.edge_types:
@@ -304,7 +305,7 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
         self,
         key: str,
         value_dict: Dict[str, Any],
-    ) -> 'HeteroData':
+    ) -> Self:
         r"""Sets the values in the dictionary :obj:`value_dict` to the
         attribute with name :obj:`key` to all node/edge types present in the
         dictionary.
@@ -324,7 +325,7 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
             self[k][key] = v
         return self
 
-    def update(self, data: 'HeteroData') -> 'HeteroData':
+    def update(self, data: Self) -> Self:
         for store in data.stores:
             for key, value in store.items():
                 self[store._key][key] = value
@@ -606,7 +607,7 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
             self._edge_store_dict[key] = out
         return out
 
-    def rename(self, name: NodeType, new_name: NodeType) -> 'HeteroData':
+    def rename(self, name: NodeType, new_name: NodeType) -> Self:
         r"""Renames the node type :obj:`name` to :obj:`new_name` in-place."""
         node_store = self._node_store_dict.pop(name)
         node_store._key = new_name
@@ -624,7 +625,7 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
 
         return self
 
-    def subgraph(self, subset_dict: Dict[NodeType, Tensor]) -> 'HeteroData':
+    def subgraph(self, subset_dict: Dict[NodeType, Tensor]) -> Self:
         r"""Returns the induced subgraph containing the node types and
         corresponding nodes in :obj:`subset_dict`.
 
@@ -718,7 +719,7 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
     def edge_subgraph(
         self,
         subset_dict: Dict[EdgeType, Tensor],
-    ) -> 'HeteroData':
+    ) -> Self:
         r"""Returns the induced subgraph given by the edge indices in
         :obj:`subset_dict` for certain edge types.
         Will currently preserve all the nodes in the graph, even if they are
@@ -742,7 +743,7 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
 
         return data
 
-    def node_type_subgraph(self, node_types: List[NodeType]) -> 'HeteroData':
+    def node_type_subgraph(self, node_types: List[NodeType]) -> Self:
         r"""Returns the subgraph induced by the given :obj:`node_types`, *i.e.*
         the returned :class:`HeteroData` object only contains the node types
         which are included in :obj:`node_types`, and only contains the edge
@@ -758,7 +759,7 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
                 del data[node_type]
         return data
 
-    def edge_type_subgraph(self, edge_types: List[EdgeType]) -> 'HeteroData':
+    def edge_type_subgraph(self, edge_types: List[EdgeType]) -> Self:
         r"""Returns the subgraph induced by the given :obj:`edge_types`, *i.e.*
         the returned :class:`HeteroData` object only contains the edge types
         which are included in :obj:`edge_types`, and only contains the node
diff --git a/torch_geometric/transforms/pad.py b/torch_geometric/transforms/pad.py
index b1846cb62..6e2a433f9 100644
--- a/torch_geometric/transforms/pad.py
+++ b/torch_geometric/transforms/pad.py
@@ -1,7 +1,5 @@
-import numbers
 from abc import ABC, abstractmethod
-from collections import defaultdict
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import torch
@@ -53,8 +51,8 @@ class UniformPadding(Padding):
 @dataclass(init=False)
 class MappingPadding(Padding):
     r"""An abstract class for specifying different padding values."""
-    values: Dict[Any, Padding] = field(default_factory=dict)
-    default: UniformPadding = 0.0
+    values: Dict[Any, Padding]
+    default: UniformPadding
 
     def __init__(
         self,
@@ -74,7 +72,7 @@ class MappingPadding(Padding):
         for key, value in self.values.items():
             self.validate_key_value(key, value)
 
-    def validate_key_value(self, key: Any, value: Any):
+    def validate_key_value(self, key: Any, value: Any) -> None:
         pass
 
 
@@ -87,7 +85,7 @@ class AttrNamePadding(MappingPadding):
             attribute names not specified in :obj:`values`.
             (default: :obj:`0.0`)
     """
-    def validate_key_value(self, key: Any, value: Any):
+    def validate_key_value(self, key: Any, value: Any) -> None:
         if not isinstance(key, str):
             raise ValueError(f"Expected the attribute name '{key}' to be a "
                              f"string (got '{type(key)}')")
@@ -113,7 +111,7 @@ class NodeTypePadding(MappingPadding):
         default (int or float, optional): The padding value to use for node
             types not specified in :obj:`values`. (default: :obj:`0.0`)
     """
-    def validate_key_value(self, key: Any, value: Any):
+    def validate_key_value(self, key: Any, value: Any) -> None:
         if not isinstance(key, str):
             raise ValueError(f"Expected the node type '{key}' to be a string "
                              f"(got '{type(key)}')")
@@ -125,7 +123,7 @@ class NodeTypePadding(MappingPadding):
 
     def get_value(
         self,
-        store_type: Optional[NodeType] = None,
+        store_type: Optional[Union[NodeType, EdgeType]] = None,
         attr_name: Optional[str] = None,
     ) -> Union[int, float]:
         padding = self.values.get(store_type, self.default)
@@ -140,7 +138,7 @@ class EdgeTypePadding(MappingPadding):
         default (int or float, optional): The padding value to use for edge
             types not specified in :obj:`values`. (default: :obj:`0.0`)
     """
-    def validate_key_value(self, key: Any, value: Any):
+    def validate_key_value(self, key: Any, value: Any) -> None:
         if not isinstance(key, tuple):
             raise ValueError(f"Expected the edge type '{key}' to be a tuple "
                              f"(got '{type(key)}')")
@@ -156,13 +154,57 @@ class EdgeTypePadding(MappingPadding):
 
     def get_value(
         self,
-        store_type: Optional[EdgeType] = None,
+        store_type: Optional[Union[NodeType, EdgeType]] = None,
         attr_name: Optional[str] = None,
     ) -> Union[int, float]:
         padding = self.values.get(store_type, self.default)
         return padding.get_value(attr_name=attr_name)
 
 
+class _NumNodes:
+    def __init__(
+        self,
+        value: Union[int, Dict[NodeType, int], None],
+    ) -> None:
+        self.value = value
+
+    def get_value(self, key: Optional[NodeType] = None) -> Optional[int]:
+        if self.value is None or isinstance(self.value, int):
+            return self.value
+        assert isinstance(key, str)
+        return self.value[key]
+
+
+class _NumEdges:
+    def __init__(
+        self,
+        value: Union[int, Dict[EdgeType, int], None],
+        num_nodes: _NumNodes,
+    ) -> None:
+
+        if value is None:
+            if isinstance(num_nodes.value, int):
+                value = num_nodes.value * num_nodes.value
+            else:
+                value = {}
+
+        self.value = value
+        self.num_nodes = num_nodes
+
+    def get_value(self, key: Optional[EdgeType] = None) -> Optional[int]:
+        if self.value is None or isinstance(self.value, int):
+            return self.value
+
+        assert isinstance(key, tuple) and len(key) == 3
+        if key not in self.value:
+            num_src_nodes = self.num_nodes.get_value(key[0])
+            num_dst_nodes = self.num_nodes.get_value(key[-1])
+            assert num_src_nodes is not None and num_dst_nodes is not None
+            self.value[key] = num_src_nodes * num_dst_nodes
+
+        return self.value[key]
+
+
 @functional_transform('pad')
 class Pad(BaseTransform):
     r"""Applies padding to enforce consistent tensor shapes
@@ -283,16 +325,20 @@ class Pad(BaseTransform):
         add_pad_mask: bool = False,
         exclude_keys: Optional[List[str]] = None,
     ):
-        self.max_num_nodes = self._NumNodes(max_num_nodes)
-        self.max_num_edges = self._NumEdges(max_num_edges, self.max_num_nodes)
+        self.max_num_nodes = _NumNodes(max_num_nodes)
+        self.max_num_edges = _NumEdges(max_num_edges, self.max_num_nodes)
 
-        self.node_pad = node_pad_value
-        if not isinstance(self.node_pad, Padding):
-            self.node_pad = UniformPadding(self.node_pad)
+        self.node_pad: Padding
+        if not isinstance(node_pad_value, Padding):
+            self.node_pad = UniformPadding(node_pad_value)
+        else:
+            self.node_pad = node_pad_value
 
-        self.edge_pad = edge_pad_value
-        if not isinstance(self.edge_pad, Padding):
-            self.edge_pad = UniformPadding(self.edge_pad)
+        self.edge_pad: Padding
+        if not isinstance(edge_pad_value, Padding):
+            self.edge_pad = UniformPadding(edge_pad_value)
+        else:
+            self.edge_pad = edge_pad_value
 
         self.node_additional_attrs_pad = {
             key: mask_pad_value
@@ -302,74 +348,6 @@ class Pad(BaseTransform):
         self.add_pad_mask = add_pad_mask
         self.exclude_keys = set(exclude_keys or [])
 
-    class _IntOrDict(ABC):
-        def __init__(self, value):
-            self.value = value
-            self.is_number = isinstance(value, numbers.Number)
-
-        @abstractmethod
-        def get_value(self, key: Optional[Any] = None):
-            pass
-
-        def is_none(self):
-            return self.value is None
-
-    class _NumNodes(_IntOrDict):
-        def __init__(self, value):
-            assert isinstance(value, (int, dict)), \
-                f'Parameter `max_num_nodes` must be of type int or dict ' \
-                f'but is {type(value)}.'
-            super().__init__(value)
-
-        def get_value(self, key: Optional[NodeType] = None) -> int:
-            if self.is_number or self.value is None:
-                # Homodata case.
-                return self.value
-
-            assert isinstance(key, str)
-            assert key in self.value.keys(), \
-                f'The number of {key} nodes was not specified for padding.'
-            return self.value[key]
-
-    class _NumEdges(_IntOrDict):
-        def __init__(self, value: Union[int, Dict[EdgeType, int], None],
-                     num_nodes: '_NumNodes'):  # noqa: F821
-            assert value is None or isinstance(value, (int, dict)), \
-                f'If provided, parameter `max_num_edges` must be of type ' \
-                f'int or dict but is {type(value)}.'
-
-            if value is not None:
-                if isinstance(value, int):
-                    num_edges = value
-                else:
-                    num_edges = defaultdict(lambda: defaultdict(int))
-                    for k, v in value.items():
-                        src_node, edge_type, dst_node = k
-                        num_edges[src_node, dst_node][edge_type] = v
-            elif num_nodes.is_number:
-                num_edges = num_nodes.get_value() * num_nodes.get_value()
-            else:
-                num_edges = defaultdict(lambda: defaultdict(int))
-
-            self.num_nodes = num_nodes
-            super().__init__(num_edges)
-
-        def get_value(self, key: Optional[EdgeType] = None) -> int:
-            if self.is_number or self.value is None:
-                # Homodata case.
-                return self.value
-
-            assert isinstance(key, tuple) and len(key) == 3
-            src_v, edge_type, dst_v = key
-            if (src_v, dst_v) in self.value.keys():
-                if edge_type in self.value[src_v, dst_v]:
-                    return self.value[src_v, dst_v][edge_type]
-
-            max_num_edges = self.num_nodes.get_value(
-                src_v) * self.num_nodes.get_value(dst_v)
-            self.value[src_v, dst_v][edge_type] = max_num_edges
-            return self.value[src_v, dst_v][edge_type]
-
     def __should_pad_node_attr(self, attr_name: str) -> bool:
         if attr_name in self.node_additional_attrs_pad:
             return True
@@ -378,7 +356,7 @@ class Pad(BaseTransform):
         return False
 
     def __should_pad_edge_attr(self, attr_name: str) -> bool:
-        if self.max_num_edges.is_none():
+        if self.max_num_edges.value is None:
             return False
         if attr_name == 'edge_index':
             return True
@@ -387,15 +365,19 @@ class Pad(BaseTransform):
         return False
 
     def __get_node_padding(
-            self, attr_name: str,
-            node_type: Optional[NodeType] = None) -> Union[int, float]:
+        self,
+        attr_name: str,
+        node_type: Optional[NodeType] = None,
+    ) -> Union[int, float]:
         if attr_name in self.node_additional_attrs_pad:
             return self.node_additional_attrs_pad[attr_name]
         return self.node_pad.get_value(node_type, attr_name)
 
     def __get_edge_padding(
-            self, attr_name: str,
-            edge_type: Optional[EdgeType] = None) -> Union[int, float]:
+        self,
+        attr_name: str,
+        edge_type: Optional[EdgeType] = None,
+    ) -> Union[int, float]:
         return self.edge_pad.get_value(edge_type, attr_name)
 
     def forward(
@@ -407,11 +389,13 @@ class Pad(BaseTransform):
             assert isinstance(self.node_pad, (UniformPadding, AttrNamePadding))
             assert isinstance(self.edge_pad, (UniformPadding, AttrNamePadding))
 
-            for store in data.stores:
-                for key in self.exclude_keys:
-                    del store[key]
-                self.__pad_edge_store(store, data.__cat_dim__, data.num_nodes)
-                self.__pad_node_store(store, data.__cat_dim__)
+            for key in self.exclude_keys:
+                del data[key]
+
+            num_nodes = data.num_nodes
+            assert num_nodes is not None
+            self.__pad_edge_store(data._store, data.__cat_dim__, num_nodes)
+            self.__pad_node_store(data._store, data.__cat_dim__)
             data.num_nodes = self.max_num_nodes.get_value()
         else:
             assert isinstance(
@@ -421,34 +405,42 @@ class Pad(BaseTransform):
                 self.edge_pad,
                 (UniformPadding, AttrNamePadding, EdgeTypePadding))
 
-            for edge_type, store in data.edge_items():
+            for edge_type, edge_store in data.edge_items():
                 for key in self.exclude_keys:
-                    del store[key]
+                    del edge_store[key]
 
                 src_node_type, _, dst_node_type = edge_type
-                self.__pad_edge_store(store, data.__cat_dim__,
-                                      (data[src_node_type].num_nodes,
-                                       data[dst_node_type].num_nodes),
+                num_src_nodes = data[src_node_type].num_nodes
+                num_dst_nodes = data[dst_node_type].num_nodes
+                assert num_src_nodes is not None and num_dst_nodes is not None
+                self.__pad_edge_store(edge_store, data.__cat_dim__,
+                                      (num_src_nodes, num_dst_nodes),
                                       edge_type)
 
-            for node_type, store in data.node_items():
+            for node_type, node_store in data.node_items():
                 for key in self.exclude_keys:
-                    del store[key]
-                self.__pad_node_store(store, data.__cat_dim__, node_type)
+                    del node_store[key]
+                self.__pad_node_store(node_store, data.__cat_dim__, node_type)
                 data[node_type].num_nodes = self.max_num_nodes.get_value(
                     node_type)
 
         return data
 
-    def __pad_node_store(self, store: NodeStorage, get_dim_fn: Callable,
-                         node_type: Optional[str] = None):
+    def __pad_node_store(
+        self,
+        store: NodeStorage,
+        get_dim_fn: Callable,
+        node_type: Optional[NodeType] = None,
+    ) -> None:
 
-        attrs_to_pad = [key for key in store if store.is_node_attr(key)]
+        attrs_to_pad = [key for key in store.keys() if store.is_node_attr(key)]
 
-        if not attrs_to_pad:
+        if len(attrs_to_pad) == 0:
             return
 
         num_target_nodes = self.max_num_nodes.get_value(node_type)
+        assert num_target_nodes is not None
+        assert store.num_nodes is not None
         assert num_target_nodes >= store.num_nodes, \
             f'The number of nodes after padding ({num_target_nodes}) cannot ' \
             f'be lower than the number of nodes in the data object ' \
@@ -467,15 +459,21 @@ class Pad(BaseTransform):
             store[attr_name] = self._pad_tensor_dim(attr, dim, num_pad_nodes,
                                                     pad_value)
 
-    def __pad_edge_store(self, store: EdgeStorage, get_dim_fn: Callable,
-                         num_nodes: Union[int, Tuple[int, int]],
-                         edge_type: Optional[str] = None):
+    def __pad_edge_store(
+        self,
+        store: EdgeStorage,
+        get_dim_fn: Callable,
+        num_nodes: Union[int, Tuple[int, int]],
+        edge_type: Optional[EdgeType] = None,
+    ) -> None:
+
         attrs_to_pad = set(
             attr for attr in store.keys()
             if store.is_edge_attr(attr) and self.__should_pad_edge_attr(attr))
         if not attrs_to_pad:
             return
         num_target_edges = self.max_num_edges.get_value(edge_type)
+        assert num_target_edges is not None
         assert num_target_edges >= store.num_edges, \
             f'The number of edges after padding ({num_target_edges}) cannot ' \
             f'be lower than the number of edges in the data object ' \
diff --git a/torch_geometric/transforms/rooted_subgraph.py b/torch_geometric/transforms/rooted_subgraph.py
index d06c0cb93..5ad6563b3 100644
--- a/torch_geometric/transforms/rooted_subgraph.py
+++ b/torch_geometric/transforms/rooted_subgraph.py
@@ -28,7 +28,7 @@ class RootedSubgraphData(Data):
     * :obj:`e_sub_batch` (Tensor): The batch vector to distinguish edges across
       different subgraphs.
     """
-    def __inc__(self, key, value, *args, **kwargs) -> Any:
+    def __inc__(self, key: str, value: Any, *args: Any, **kwargs: Any) -> Any:
         if key == 'sub_edge_index':
             return self.n_id.size(0)
         if key in ['n_sub_batch', 'e_sub_batch']:
@@ -36,6 +36,7 @@ class RootedSubgraphData(Data):
         elif key == 'n_id':
             return self.num_nodes
         elif key == 'e_id':
+            assert self.edge_index is not None
             return self.edge_index.size(1)
         return super().__inc__(key, value, *args, **kwargs)
 
@@ -81,13 +82,17 @@ class RootedSubgraph(BaseTransform, ABC):
         n_mask: Tensor,
     ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
 
+        assert data.edge_index is not None
+        num_nodes = data.num_nodes
+        assert num_nodes is not None
+
         n_sub_batch, n_id = n_mask.nonzero().t()
         e_mask = n_mask[:, data.edge_index[0]] & n_mask[:, data.edge_index[1]]
         e_sub_batch, e_id = e_mask.nonzero().t()
 
         sub_edge_index = data.edge_index[:, e_id]
         arange = torch.arange(n_id.size(0), device=data.edge_index.device)
-        node_map = data.edge_index.new_ones(data.num_nodes, data.num_nodes)
+        node_map = data.edge_index.new_ones(num_nodes, num_nodes)
         node_map[n_sub_batch, n_id] = arange
         sub_edge_index += (arange * data.num_nodes)[e_sub_batch]
         sub_edge_index = node_map.view(-1)[sub_edge_index]
@@ -109,7 +114,7 @@ class RootedEgoNets(RootedSubgraph):
     Args:
         num_hops (int): the number of hops :math:`k`.
     """
-    def __init__(self, num_hops: int):
+    def __init__(self, num_hops: int) -> None:
         super().__init__()
         self.num_hops = num_hops
 
@@ -118,8 +123,12 @@ class RootedEgoNets(RootedSubgraph):
         data: Data,
     ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
 
+        assert data.edge_index is not None
+        num_nodes = data.num_nodes
+        assert num_nodes is not None
+
         adj_t = to_torch_csc_tensor(data.edge_index, size=data.size()).t()
-        n_mask = torch.eye(data.num_nodes, device=data.edge_index.device)
+        n_mask = torch.eye(num_nodes, device=data.edge_index.device)
         for _ in range(self.num_hops):
             n_mask += adj_t @ n_mask
 
@@ -150,13 +159,17 @@ class RootedRWSubgraph(RootedSubgraph):
     ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
         from torch_cluster import random_walk
 
-        start = torch.arange(data.num_nodes, device=data.edge_index.device)
+        assert data.edge_index is not None
+        num_nodes = data.num_nodes
+        assert num_nodes is not None
+
+        start = torch.arange(num_nodes, device=data.edge_index.device)
         start = start.view(-1, 1).repeat(1, self.repeat).view(-1)
         walk = random_walk(data.edge_index[0], data.edge_index[1], start,
                            self.walk_length, num_nodes=data.num_nodes)
 
-        n_mask = torch.zeros((data.num_nodes, data.num_nodes),
-                             dtype=torch.bool, device=walk.device)
+        n_mask = torch.zeros((num_nodes, num_nodes), dtype=torch.bool,
+                             device=walk.device)
         start = start.view(-1, 1).repeat(1, (self.walk_length + 1)).view(-1)
         n_mask[start, walk.view(-1)] = True
 


Test patch:
diff --git a/test/transforms/test_pad.py b/test/transforms/test_pad.py
index 944ed9e56..fac3b3513 100644
--- a/test/transforms/test_pad.py
+++ b/test/transforms/test_pad.py
@@ -449,7 +449,7 @@ def test_pad_num_nodes_not_complete():
     data = fake_hetero_data(node_types=2, edge_types=1)
     transform = Pad(max_num_nodes={'v0': 100})
 
-    with pytest.raises(AssertionError, match="The number of v1 nodes"):
+    with pytest.raises(KeyError):
         transform(data)