Merge branch 'master' into type_hints/ToDense

Author: rusty1s

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), [#5668](https://github.com/pyg-team/pytorch_geometric/pull/5668))
+- 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))

test/utils/test_augmentation.py

@@ -54,20 +54,21 @@ def test_mask_feature():
 
     torch.manual_seed(7)
     out = mask_feature(x, mode='all')
-    assert out[0].tolist() == [[1.0, 0.0, 0.0, 4.0], [5.0, 6.0, 7.0, 0.0],
-                               [9.0, 10.0, 0.0, 12.0]]
+    assert out[0].tolist() == [[0.0, 2.0, 3.0, 0.0], [5.0, 0.0, 0.0, 8.0],
+                               [0.0, 0.0, 0.0, 0.0]]
 
-    assert out[1].tolist() == [[True, False, False, True],
-                               [True, True, True, False],
-                               [True, True, False, True]]
+    assert out[1].tolist() == [[False, True, True, False],
+                               [True, False, False, True],
+                               [False, False, False, False]]
 
     torch.manual_seed(7)
     out = mask_feature(x, mode='all', fill_value=-1)
-    assert out[0].tolist() == [[1.0, -1., -1., 4.0], [5.0, 6.0, 7.0, -1.],
-                               [9.0, 10.0, -1., 12.0]]
-    assert out[1].tolist() == [[True, False, False, True],
-                               [True, True, True, False],
-                               [True, True, False, True]]
+    assert out[0].tolist() == [[-1.0, 2.0, 3.0, -1.0], [5.0, -1.0, -1.0, 8.0],
+                               [-1.0, -1.0, -1.0, -1.0]]
+
+    assert out[1].tolist() == [[False, True, True, False],
+                               [True, False, False, True],
+                               [False, False, False, False]]
 
 
 def test_add_random_edge():

test/utils/test_isolated.py

@@ -28,6 +28,12 @@ def test_remove_isolated_nodes():
     assert out.tolist() == [[0, 1, 0], [1, 0, 0]]
     assert mask.tolist() == [1, 1]
 
+    if is_full_test():
+        jit = torch.jit.script(remove_isolated_nodes)
+        out, _, mask = jit(edge_index)
+        assert out.tolist() == [[0, 1, 0], [1, 0, 0]]
+        assert mask.tolist() == [1, 1]
+
     out, _, mask = remove_isolated_nodes(edge_index, num_nodes=3)
     assert out.tolist() == [[0, 1, 0], [1, 0, 0]]
     assert mask.tolist() == [1, 1, 0]

test/utils/test_normalized_cut.py

@@ -1,5 +1,6 @@
 import torch
 
+from torch_geometric.testing import is_full_test
 from torch_geometric.utils import normalized_cut
 
 
@@ -11,3 +12,8 @@ def test_normalized_cut():
 
     output = normalized_cut(torch.stack([row, col], dim=0), edge_attr)
     assert output.tolist() == expected_output
+
+    if is_full_test():
+        jit = torch.jit.script(normalized_cut)
+        output = jit(torch.stack([row, col], dim=0), edge_attr)
+        assert output.tolist() == expected_output

torch_geometric/datasets/icews.py

@@ -1,26 +1,33 @@
+from typing import Callable, List, Optional
+
 import torch
 
 from torch_geometric.data import Data, InMemoryDataset, download_url
 from torch_geometric.io import read_txt_array
 
 
 class EventDataset(InMemoryDataset):
-    def __init__(self, root, transform=None, pre_transform=None,
-                 pre_filter=None):
+    def __init__(
+        self,
+        root: str,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        pre_filter: Optional[Callable] = None,
+    ):
         super().__init__(root, transform, pre_transform, pre_filter)
 
     @property
-    def num_nodes(self):
+    def num_nodes(self) -> int:
         raise NotImplementedError
 
     @property
-    def num_rels(self):
+    def num_rels(self) -> int:
         raise NotImplementedError
 
-    def process_events(self):
+    def process_events(self) -> int:
         raise NotImplementedError
 
-    def process(self):
+    def process(self) -> List[Data]:
         events = self.process_events()
         events = events - events.min(dim=0, keepdim=True)[0]
 
@@ -64,34 +71,40 @@ class ICEWS18(EventDataset):
     url = 'https://github.com/INK-USC/RE-Net/raw/master/data/ICEWS18'
     splits = [0, 373018, 419013, 468558]  # Train/Val/Test splits.
 
-    def __init__(self, root, split='train', transform=None, pre_transform=None,
-                 pre_filter=None):
+    def __init__(
+        self,
+        root: str,
+        split: str = 'train',
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        pre_filter: Optional[Callable] = None,
+    ):
         assert split in ['train', 'val', 'test']
         super().__init__(root, transform, pre_transform, pre_filter)
         idx = self.processed_file_names.index(f'{split}.pt')
         self.data, self.slices = torch.load(self.processed_paths[idx])
 
     @property
-    def num_nodes(self):
+    def num_nodes(self) -> int:
         return 23033
 
     @property
-    def num_rels(self):
+    def num_rels(self) -> int:
         return 256
 
     @property
-    def raw_file_names(self):
+    def raw_file_names(self) -> List[str]:
         return [f'{name}.txt' for name in ['train', 'valid', 'test']]
 
     @property
-    def processed_file_names(self):
+    def processed_file_names(self) -> List[str]:
         return ['train.pt', 'val.pt', 'test.pt']
 
     def download(self):
         for filename in self.raw_file_names:
             download_url(f'{self.url}/{filename}', self.raw_dir)
 
-    def process_events(self):
+    def process_events(self) -> torch.Tensor:
         events = []
         for path in self.raw_paths:
             data = read_txt_array(path, sep='\t', end=4, dtype=torch.long)

torch_geometric/transforms/one_hot_degree.py

@@ -1,6 +1,7 @@
 import torch
 import torch.nn.functional as F
 
+from torch_geometric.data import Data
 from torch_geometric.data.datapipes import functional_transform
 from torch_geometric.transforms import BaseTransform
 from torch_geometric.utils import degree
@@ -19,12 +20,17 @@ class OneHotDegree(BaseTransform):
         cat (bool, optional): Concat node degrees to node features instead
             of replacing them. (default: :obj:`True`)
     """
-    def __init__(self, max_degree, in_degree=False, cat=True):
+    def __init__(
+        self,
+        max_degree: int,
+        in_degree: bool = False,
+        cat: bool = True,
+    ):
         self.max_degree = max_degree
         self.in_degree = in_degree
         self.cat = cat
 
-    def __call__(self, data):
+    def __call__(self, data: Data) -> Data:
         idx, x = data.edge_index[1 if self.in_degree else 0], data.x
         deg = degree(idx, data.num_nodes, dtype=torch.long)
         deg = F.one_hot(deg, num_classes=self.max_degree + 1).to(torch.float)

torch_geometric/utils/augmentation.py

@@ -143,7 +143,7 @@ def mask_feature(x: Tensor, p: float = 0.5, mode: str = 'col',
         mask = torch.rand(x.size(1), device=x.device) >= p
         mask = mask.view(1, -1)
     else:
-        mask = x.bernoulli(1 - p).to(torch.bool)
+        mask = torch.randn_like(x) >= p
 
     x = x.masked_fill(~mask, fill_value)
     return x, mask

torch_geometric/utils/isolated.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, Tuple
 
 import torch
 from torch import Tensor
@@ -37,7 +37,11 @@ def contains_isolated_nodes(
     return torch.unique(edge_index.view(-1)).numel() < num_nodes
 
 
-def remove_isolated_nodes(edge_index, edge_attr=None, num_nodes=None):
+def remove_isolated_nodes(
+    edge_index: Tensor,
+    edge_attr: Optional[Tensor] = None,
+    num_nodes: Optional[int] = None,
+) -> Tuple[Tensor, Optional[Tensor], Tensor]:
     r"""Removes the isolated nodes from the graph given by :attr:`edge_index`
     with optional edge attributes :attr:`edge_attr`.
     In addition, returns a mask of shape :obj:`[num_nodes]` to manually filter

torch_geometric/utils/normalized_cut.py

@@ -1,9 +1,15 @@
 from typing import Optional
 
+from torch import Tensor
+
 from torch_geometric.utils import degree
 
 
-def normalized_cut(edge_index, edge_attr, num_nodes: Optional[int] = None):
+def normalized_cut(
+    edge_index: Tensor,
+    edge_attr: Tensor,
+    num_nodes: Optional[int] = None,
+) -> Tensor:
     r"""Computes the normalized cut :math:`\mathbf{e}_{i,j} \cdot
     \left( \frac{1}{\deg(i)} + \frac{1}{\deg(j)} \right)` of a weighted graph
     given by edge indices and edge attributes.