[TypeHints] ToDense (#5668)

Co-authored-by: rusty1s <[email protected]>

Author: sebastian-montero

CHANGELOG.md

@@ -40,7 +40,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Support `in_channels` with `tuple` in `GENConv` for bipartite message passing ([#5627](https://github.com/pyg-team/pytorch_geometric/pull/5627), [#5641](https://github.com/pyg-team/pytorch_geometric/pull/5641))
 - Handle cases of not having enough possible negative edges in `RandomLinkSplit` ([#5642](https://github.com/pyg-team/pytorch_geometric/pull/5642))
 - Fix `RGCN+pyg-lib` for `LongTensor` input ([#5610](https://github.com/pyg-team/pytorch_geometric/pull/5610))
-- Improved type hint support ([#5603](https://github.com/pyg-team/pytorch_geometric/pull/5603), [#5659](https://github.com/pyg-team/pytorch_geometric/pull/5659), [#5665](https://github.com/pyg-team/pytorch_geometric/pull/5665), [#5666](https://github.com/pyg-team/pytorch_geometric/pull/5666), [#5667](https://github.com/pyg-team/pytorch_geometric/pull/5667))
+- Improved type hint support ([#5603](https://github.com/pyg-team/pytorch_geometric/pull/5603), [#5659](https://github.com/pyg-team/pytorch_geometric/pull/5659), [#5665](https://github.com/pyg-team/pytorch_geometric/pull/5665), [#5666](https://github.com/pyg-team/pytorch_geometric/pull/5666), [#5667](https://github.com/pyg-team/pytorch_geometric/pull/5667), [#5668](https://github.com/pyg-team/pytorch_geometric/pull/5668))
 - Avoid modifying `mode_kwargs` in `MultiAggregation` ([#5601](https://github.com/pyg-team/pytorch_geometric/pull/5601))
 - Changed `BatchNorm` to allow for batches of size one during training ([#5530](https://github.com/pyg-team/pytorch_geometric/pull/5530), [#5614](https://github.com/pyg-team/pytorch_geometric/pull/5614))
 - Integrated better temporal sampling support by requiring that local neighborhoods are sorted according to time ([#5516](https://github.com/pyg-team/pytorch_geometric/issues/5516), [#5602](https://github.com/pyg-team/pytorch_geometric/issues/5602))

torch_geometric/transforms/to_dense.py

@@ -1,5 +1,8 @@
+from typing import Optional
+
 import torch
 
+from torch_geometric.data import Data
 from torch_geometric.data.datapipes import functional_transform
 from torch_geometric.transforms import BaseTransform
 
@@ -10,13 +13,14 @@ class ToDense(BaseTransform):
     shape :obj:`[num_nodes, num_nodes, *]` (functional name: :obj:`to_dense`).
 
     Args:
-        num_nodes (int): The number of nodes. If set to :obj:`None`, the number
-            of nodes will get automatically inferred. (default: :obj:`None`)
+        num_nodes (int, optional): The number of nodes. If set to :obj:`None`,
+            the number of nodes will get automatically inferred.
+            (default: :obj:`None`)
     """
-    def __init__(self, num_nodes=None):
+    def __init__(self, num_nodes: Optional[int] = None):
         self.num_nodes = num_nodes
 
-    def __call__(self, data):
+    def __call__(self, data: Data) -> Data:
         assert data.edge_index is not None
 
         orig_num_nodes = data.num_nodes