Merge pull request #1 from pyg-team/master

Update to latest

docs/source/tutorial/multi_node_multi_gpu_vanilla.rst

@@ -61,9 +61,8 @@ Finally, to submit the :obj:`*.sbatch` file itself into the work queue, use the
 Using a cluster configured with pyxis-containers
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-If your cluster supports the :obj:`pyxis` plugin developed by NVIDIA, you can use a ready-to-use :pyg:`PyG` container that is updated each month with the latest from NVIDIA and :pyg:`PyG`.
-Currently it is not yet publically available, but you can sign up for early access `here <https://developer.nvidia.com/pyg-container-early-access>`_.
-The container should set up all necessary environment variables from which you can now directly run the example using :obj:`srun` from your command prompt:
+If your cluster supports the :obj:`pyxis` plugin developed by NVIDIA, you can use a ready-to-use :pyg:`PyG` container that is updated each month with the latest from NVIDIA and :pyg:`PyG`, see `here <https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pyg>`_ for more information.
+The container sets up all necessary environment variables from which you can now directly run the example using :obj:`srun` from your command prompt:
 
 .. code-block:: console
 

test/nn/models/test_basic_gnn.py

@@ -246,7 +246,7 @@ def test_packaging():
     path = osp.join(torch.hub._get_torch_home(), 'pyg_test_package.pt')
     with torch.package.PackageExporter(path) as pe:
         pe.extern('torch_geometric.nn.**')
-        pe.extern('torch_geometric.utils.trim_to_layer')
+        pe.extern('torch_geometric.utils._trim_to_layer')
         pe.extern('_operator')
         pe.save_pickle('models', 'model.pkl', model)
 

test/utils/test_negative_sampling.py

@@ -9,7 +9,7 @@
     structured_negative_sampling_feasible,
     to_undirected,
 )
-from torch_geometric.utils.negative_sampling import (
+from torch_geometric.utils._negative_sampling import (
     edge_index_to_vector,
     vector_to_edge_index,
 )

test/utils/test_scatter.py

@@ -6,7 +6,7 @@
 from torch_geometric.profile import benchmark
 from torch_geometric.testing import disableExtensions, withCUDA, withPackage
 from torch_geometric.utils import group_argsort, scatter
-from torch_geometric.utils.scatter import scatter_argmax
+from torch_geometric.utils._scatter import scatter_argmax
 
 
 @withPackage('torch>=1.12.0')

test/utils/test_trim_to_layer.py

@@ -10,7 +10,7 @@
 from torch_geometric.testing import withPackage
 from torch_geometric.typing import SparseTensor
 from torch_geometric.utils import trim_to_layer
-from torch_geometric.utils.trim_to_layer import trim_sparse_tensor
+from torch_geometric.utils._trim_to_layer import trim_sparse_tensor
 
 
 @withPackage('torch_sparse')

torch_geometric/contrib/explain/pgm_explainer.py

@@ -9,7 +9,7 @@
 from torch_geometric.explain.config import ModelMode, ModelTaskLevel
 from torch_geometric.explain.explanation import Explanation
 from torch_geometric.utils import k_hop_subgraph
-from torch_geometric.utils.subgraph import get_num_hops
+from torch_geometric.utils._subgraph import get_num_hops
 
 
 class PGMExplainer(ExplainerAlgorithm):

torch_geometric/data/hypergraph_data.py

@@ -8,7 +8,7 @@
 from torch_geometric.data import Data
 from torch_geometric.typing import EdgeType, NodeType, OptTensor
 from torch_geometric.utils import select
-from torch_geometric.utils.subgraph import hyper_subgraph
+from torch_geometric.utils._subgraph import hyper_subgraph
 
 
 class HyperGraphData(Data):

torch_geometric/nn/models/basic_gnn.py

@@ -26,7 +26,7 @@
     normalization_resolver,
 )
 from torch_geometric.typing import Adj, OptTensor, SparseTensor
-from torch_geometric.utils.trim_to_layer import TrimToLayer
+from torch_geometric.utils._trim_to_layer import TrimToLayer
 
 
 class BasicGNN(torch.nn.Module):

torch_geometric/nn/models/tgn.py

@@ -7,7 +7,7 @@
 
 from torch_geometric.nn.inits import zeros
 from torch_geometric.utils import scatter
-from torch_geometric.utils.scatter import scatter_argmax
+from torch_geometric.utils._scatter import scatter_argmax
 
 TGNMessageStoreType = Dict[int, Tuple[Tensor, Tensor, Tensor, Tensor]]
 

torch_geometric/utils/__init__.py

@@ -2,41 +2,41 @@
 
 import copy
 
-from .scatter import scatter, group_argsort
-from .segment import segment
-from .sort import index_sort
+from ._scatter import scatter, group_argsort
+from ._segment import segment
+from ._index_sort import index_sort
 from .functions import cumsum
-from .degree import degree
-from .softmax import softmax
-from .sort_edge_index import sort_edge_index
-from .lexsort import lexsort
-from .coalesce import coalesce
+from ._degree import degree
+from ._softmax import softmax
+from ._sort_edge_index import sort_edge_index
+from ._lexsort import lexsort
+from ._coalesce import coalesce
 from .undirected import is_undirected, to_undirected
 from .loop import (contains_self_loops, remove_self_loops,
                    segregate_self_loops, add_self_loops,
                    add_remaining_self_loops, get_self_loop_attr)
 from .isolated import contains_isolated_nodes, remove_isolated_nodes
-from .subgraph import (get_num_hops, subgraph, k_hop_subgraph,
-                       bipartite_subgraph)
+from ._subgraph import (get_num_hops, subgraph, k_hop_subgraph,
+                        bipartite_subgraph)
 from .dropout import dropout_adj, dropout_node, dropout_edge, dropout_path
-from .homophily import homophily
-from .assortativity import assortativity
-from .get_laplacian import get_laplacian
-from .get_mesh_laplacian import get_mesh_laplacian
+from ._homophily import homophily
+from ._assortativity import assortativity
+from .laplacian import get_laplacian
+from .mesh_laplacian import get_mesh_laplacian
 from .mask import mask_select, index_to_mask, mask_to_index
-from .select import select, narrow
-from .to_dense_batch import to_dense_batch
-from .to_dense_adj import to_dense_adj
+from ._select import select, narrow
+from ._to_dense_batch import to_dense_batch
+from ._to_dense_adj import to_dense_adj
 from .nested import to_nested_tensor, from_nested_tensor
 from .sparse import (dense_to_sparse, is_sparse, is_torch_sparse_tensor,
                      to_torch_coo_tensor, to_torch_csr_tensor,
                      to_torch_csc_tensor, to_torch_sparse_tensor,
                      to_edge_index)
-from .spmm import spmm
-from .unbatch import unbatch, unbatch_edge_index
-from .one_hot import one_hot
-from .normalized_cut import normalized_cut
-from .grid import grid
+from ._spmm import spmm
+from ._unbatch import unbatch, unbatch_edge_index
+from ._one_hot import one_hot
+from ._normalized_cut import normalized_cut
+from ._grid import grid
 from .geodesic import geodesic_distance
 from .convert import to_scipy_sparse_matrix, from_scipy_sparse_matrix
 from .convert import to_networkx, from_networkx
@@ -47,15 +47,15 @@
 from .smiles import from_smiles, to_smiles
 from .random import (erdos_renyi_graph, stochastic_blockmodel_graph,
                      barabasi_albert_graph)
-from .negative_sampling import (negative_sampling, batched_negative_sampling,
-                                structured_negative_sampling,
-                                structured_negative_sampling_feasible)
+from ._negative_sampling import (negative_sampling, batched_negative_sampling,
+                                 structured_negative_sampling,
+                                 structured_negative_sampling_feasible)
 from .augmentation import shuffle_node, mask_feature, add_random_edge
-from .tree_decomposition import tree_decomposition
+from ._tree_decomposition import tree_decomposition
 from .embedding import get_embeddings
-from .trim_to_layer import trim_to_layer
+from ._trim_to_layer import trim_to_layer
 from .ppr import get_ppr
-from .train_test_split_edges import train_test_split_edges
+from ._train_test_split_edges import train_test_split_edges
 
 __all__ = [
     'scatter',

torch_geometric/utils/assortativity.py → torch_geometric/utils/_assortativity.py

@@ -3,7 +3,7 @@
 
 from torch_geometric.typing import Adj, SparseTensor
 from torch_geometric.utils import coalesce, degree
-from torch_geometric.utils.to_dense_adj import to_dense_adj
+from torch_geometric.utils._to_dense_adj import to_dense_adj
 
 
 def assortativity(edge_index: Adj) -> float:

torch_geometric/utils/coalesce.py → torch_geometric/utils/_coalesce.py

@@ -1,3 +1,4 @@
+import typing
 from typing import List, Optional, Tuple, Union
 
 import torch
@@ -7,35 +8,67 @@
 from torch_geometric.utils import index_sort, scatter
 from torch_geometric.utils.num_nodes import maybe_num_nodes
 
+if typing.TYPE_CHECKING:
+    from typing import overload
+else:
+    from torch.jit import _overload as overload
+
 MISSING = '???'
 
 
-@torch.jit._overload
+@overload
+def coalesce(
+    edge_index: Tensor,
+    edge_attr: str = MISSING,
+    num_nodes: Optional[int] = None,
+    reduce: str = 'sum',
+    is_sorted: bool = False,
+    sort_by_row: bool = True,
+) -> Tensor:
+    pass
+
+
+@overload
 def coalesce(  # noqa: F811
-        edge_index, edge_attr, num_nodes, reduce, is_sorted, sort_by_row):
-    # type: (Tensor, str, Optional[int], str, bool, bool) -> Tensor
+    edge_index: Tensor,
+    edge_attr: Tensor,
+    num_nodes: Optional[int] = None,
+    reduce: str = 'sum',
+    is_sorted: bool = False,
+    sort_by_row: bool = True,
+) -> Tuple[Tensor, Tensor]:
     pass
 
 
-@torch.jit._overload
+@overload
 def coalesce(  # noqa: F811
-        edge_index, edge_attr, num_nodes, reduce, is_sorted, sort_by_row):
-    # type: (Tensor, Optional[Tensor], Optional[int], str, bool, bool) -> Tuple[Tensor, Optional[Tensor]]  # noqa
+    edge_index: Tensor,
+    edge_attr: OptTensor,
+    num_nodes: Optional[int] = None,
+    reduce: str = 'sum',
+    is_sorted: bool = False,
+    sort_by_row: bool = True,
+) -> Tuple[Tensor, OptTensor]:
     pass
 
 
-@torch.jit._overload
+@overload
 def coalesce(  # noqa: F811
-        edge_index, edge_attr, num_nodes, reduce, is_sorted, sort_by_row):
-    # type: (Tensor, List[Tensor], Optional[int], str, bool, bool) -> Tuple[Tensor, List[Tensor]]  # noqa
+    edge_index: Tensor,
+    edge_attr: List[Tensor],
+    num_nodes: Optional[int] = None,
+    reduce: str = 'sum',
+    is_sorted: bool = False,
+    sort_by_row: bool = True,
+) -> Tuple[Tensor, List[Tensor]]:
     pass
 
 
 def coalesce(  # noqa: F811
     edge_index: Tensor,
     edge_attr: Union[OptTensor, List[Tensor], str] = MISSING,
     num_nodes: Optional[int] = None,
-    reduce: str = 'add',
+    reduce: str = 'sum',
     is_sorted: bool = False,
     sort_by_row: bool = True,
 ) -> Union[Tensor, Tuple[Tensor, OptTensor], Tuple[Tensor, List[Tensor]]]:
@@ -52,8 +85,8 @@ def coalesce(  # noqa: F811
         num_nodes (int, optional): The number of nodes, *i.e.*
             :obj:`max_val + 1` of :attr:`edge_index`. (default: :obj:`None`)
         reduce (str, optional): The reduce operation to use for merging edge
-            features (:obj:`"add"`, :obj:`"mean"`, :obj:`"min"`, :obj:`"max"`,
-            :obj:`"mul"`, :obj:`"any"`). (default: :obj:`"add"`)
+            features (:obj:`"sum"`, :obj:`"mean"`, :obj:`"min"`, :obj:`"max"`,
+            :obj:`"mul"`, :obj:`"any"`). (default: :obj:`"sum"`)
         is_sorted (bool, optional): If set to :obj:`True`, will expect
             :obj:`edge_index` to be already sorted row-wise.
         sort_by_row (bool, optional): If set to :obj:`False`, will sort
@@ -117,7 +150,9 @@ def coalesce(  # noqa: F811
 
     # Only perform expensive merging in case there exists duplicates:
     if mask.all():
-        if edge_attr is None or isinstance(edge_attr, (Tensor, list, tuple)):
+        if edge_attr is None or isinstance(edge_attr, Tensor):
+            return edge_index, edge_attr
+        if isinstance(edge_attr, (list, tuple)):
             return edge_index, edge_attr
         return edge_index
 

torch_geometric/utils/grid.py → torch_geometric/utils/_grid.py

@@ -3,7 +3,7 @@
 import torch
 from torch import Tensor
 
-from torch_geometric.utils.coalesce import coalesce
+from torch_geometric.utils import coalesce
 
 
 def grid(
@@ -49,8 +49,10 @@ def grid_index(
 ) -> Tensor:
 
     w = width
-    kernel = [-w - 1, -1, w - 1, -w, 0, w, -w + 1, 1, w + 1]
-    kernel = torch.tensor(kernel, device=device)
+    kernel = torch.tensor(
+        [-w - 1, -1, w - 1, -w, 0, w, -w + 1, 1, w + 1],
+        device=device,
+    )
 
     row = torch.arange(height * width, dtype=torch.long, device=device)
     row = row.view(-1, 1).repeat(1, kernel.size(0))

torch_geometric/utils/homophily.py → torch_geometric/utils/_homophily.py

@@ -1,4 +1,4 @@
-from typing import Union
+from typing import Union, overload
 
 import torch
 from torch import Tensor
@@ -7,8 +7,32 @@
 from torch_geometric.utils import degree, scatter
 
 
-def homophily(edge_index: Adj, y: Tensor, batch: OptTensor = None,
-              method: str = 'edge') -> Union[float, Tensor]:
+@overload
+def homophily(
+    edge_index: Adj,
+    y: Tensor,
+    batch: None = ...,
+    method: str = ...,
+) -> float:
+    pass
+
+
+@overload
+def homophily(
+    edge_index: Adj,
+    y: Tensor,
+    batch: Tensor,
+    method: str = ...,
+) -> Tensor:
+    pass
+
+
+def homophily(
+    edge_index: Adj,
+    y: Tensor,
+    batch: OptTensor = None,
+    method: str = 'edge',
+) -> Union[float, Tensor]:
     r"""The homophily of a graph characterizes how likely nodes with the same
     label are near each other in a graph.
 

torch_geometric/utils/negative_sampling.py → torch_geometric/utils/_negative_sampling.py

@@ -55,11 +55,16 @@ def negative_sampling(
     """
     assert method in ['sparse', 'dense']
 
-    size = num_nodes
-    bipartite = isinstance(size, (tuple, list))
-    size = maybe_num_nodes(edge_index) if size is None else size
-    size = (size, size) if not bipartite else size
-    force_undirected = False if bipartite else force_undirected
+    if num_nodes is None:
+        num_nodes = maybe_num_nodes(edge_index, num_nodes)
+
+    if isinstance(num_nodes, int):
+        size = (num_nodes, num_nodes)
+        bipartite = False
+    else:
+        size = num_nodes
+        bipartite = True
+        force_undirected = False
 
     idx, population = edge_index_to_vector(edge_index, size, bipartite,
                                            force_undirected)
@@ -95,7 +100,7 @@ def negative_sampling(
         idx = idx.to('cpu')
         for _ in range(3):  # Number of tries to sample negative indices.
             rnd = sample(population, sample_size, device='cpu')
-            mask = np.isin(rnd, idx)
+            mask = np.isin(rnd.numpy(), idx.numpy())
             if neg_idx is not None:
                 mask |= np.isin(rnd, neg_idx.to('cpu'))
             mask = torch.from_numpy(mask).to(torch.bool)

torch_geometric/utils/select.py → torch_geometric/utils/_select.py

@@ -52,7 +52,7 @@ def narrow(src: Union[Tensor, List[Any]], dim: int, start: int,
         start (int): The starting dimension.
         length (int): The distance to the ending dimension.
     """
-    if is_torch_sparse_tensor(src):
+    if isinstance(src, Tensor) and is_torch_sparse_tensor(src):
         # TODO Sparse tensors in `torch.sparse` do not yet support `narrow`.
         index = torch.arange(start, start + length, device=src.device)
         return src.index_select(dim, index)