pyg-nightly 2.7.0.dev20241009__py3-none-any.whl → 2.8.0.dev20251207__py3-none-any.whl

torch_geometric/nn/norm/layer_norm.py

@@ -30,11 +30,13 @@ class LayerNorm(torch.nn.Module):
         affine (bool, optional): If set to :obj:`True`, this module has
             learnable affine parameters :math:`\gamma` and :math:`\beta`.
             (default: :obj:`True`)
-        mode (str, optinal): The normalization mode to use for layer
+        mode (str, optional): The normalization mode to use for layer
             normalization (:obj:`"graph"` or :obj:`"node"`). If :obj:`"graph"`
             is used, each graph will be considered as an element to be
             normalized. If `"node"` is used, each node will be considered as
             an element to be normalized. (default: :obj:`"graph"`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -42,6 +44,7 @@ class LayerNorm(torch.nn.Module):
         eps: float = 1e-5,
         affine: bool = True,
         mode: str = 'graph',
+        device: Optional[torch.device] = None,
     ):
         super().__init__()
 
@@ -51,8 +54,8 @@ class LayerNorm(torch.nn.Module):
         self.mode = mode
 
         if affine:
-            self.weight = Parameter(torch.empty(in_channels))
-            self.bias = Parameter(torch.empty(in_channels))
+            self.weight = Parameter(torch.empty(in_channels, device=device))
+            self.bias = Parameter(torch.empty(in_channels, device=device))
         else:
             self.register_parameter('weight', None)
             self.register_parameter('bias', None)
@@ -108,7 +111,7 @@ class LayerNorm(torch.nn.Module):
             return F.layer_norm(x, (self.in_channels, ), self.weight,
                                 self.bias, self.eps)
 
-        raise ValueError(f"Unknow normalization mode: {self.mode}")
+        raise ValueError(f"Unknownn normalization mode: {self.mode}")
 
     def __repr__(self):
         return (f'{self.__class__.__name__}({self.in_channels}, '
@@ -130,10 +133,12 @@ class HeteroLayerNorm(torch.nn.Module):
         affine (bool, optional): If set to :obj:`True`, this module has
             learnable affine parameters :math:`\gamma` and :math:`\beta`.
             (default: :obj:`True`)
-        mode (str, optinal): The normalization mode to use for layer
+        mode (str, optional): The normalization mode to use for layer
             normalization (:obj:`"node"`). If `"node"` is used, each node will
             be considered as an element to be normalized.
             (default: :obj:`"node"`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -142,6 +147,7 @@ class HeteroLayerNorm(torch.nn.Module):
         eps: float = 1e-5,
         affine: bool = True,
         mode: str = 'node',
+        device: Optional[torch.device] = None,
     ):
         super().__init__()
         assert mode == 'node'
@@ -152,8 +158,10 @@ class HeteroLayerNorm(torch.nn.Module):
         self.affine = affine
 
         if affine:
-            self.weight = Parameter(torch.empty(num_types, in_channels))
-            self.bias = Parameter(torch.empty(num_types, in_channels))
+            self.weight = Parameter(
+                torch.empty(num_types, in_channels, device=device))
+            self.bias = Parameter(
+                torch.empty(num_types, in_channels, device=device))
         else:
             self.register_parameter('weight', None)
             self.register_parameter('bias', None)

torch_geometric/nn/norm/msg_norm.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import torch
 import torch.nn.functional as F
 from torch import Tensor
@@ -19,10 +21,14 @@ class MessageNorm(torch.nn.Module):
         learn_scale (bool, optional): If set to :obj:`True`, will learn the
             scaling factor :math:`s` of message normalization.
             (default: :obj:`False`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
-    def __init__(self, learn_scale: bool = False):
+    def __init__(self, learn_scale: bool = False,
+                 device: Optional[torch.device] = None):
         super().__init__()
-        self.scale = Parameter(torch.empty(1), requires_grad=learn_scale)
+        self.scale = Parameter(torch.empty(1, device=device),
+                               requires_grad=learn_scale)
         self.reset_parameters()
 
     def reset_parameters(self):

torch_geometric/nn/pool/__init__.py

@@ -163,8 +163,10 @@ def knn_graph(
     :rtype: :class:`torch.Tensor`
     """
     if batch is not None and x.device != batch.device:
-        warnings.warn("Input tensor 'x' and 'batch' are on different devices "
-                      "in 'knn_graph'. Performing blocking device transfer")
+        warnings.warn(
+            "Input tensor 'x' and 'batch' are on different devices "
+            "in 'knn_graph'. Performing blocking device transfer",
+            stacklevel=2)
         batch = batch.to(x.device)
 
     if not torch_geometric.typing.WITH_TORCH_CLUSTER_BATCH_SIZE:
@@ -285,8 +287,10 @@ def radius_graph(
         inputs to GPU before proceeding.
     """
     if batch is not None and x.device != batch.device:
-        warnings.warn("Input tensor 'x' and 'batch' are on different devices "
-                      "in 'radius_graph'. Performing blocking device transfer")
+        warnings.warn(
+            "Input tensor 'x' and 'batch' are on different devices "
+            "in 'radius_graph'. Performing blocking device transfer",
+            stacklevel=2)
         batch = batch.to(x.device)
 
     if not torch_geometric.typing.WITH_TORCH_CLUSTER_BATCH_SIZE:

torch_geometric/nn/pool/cluster_pool.py

@@ -20,8 +20,7 @@ class UnpoolInfo(NamedTuple):
 
 class ClusterPooling(torch.nn.Module):
     r"""The cluster pooling operator from the `"Edge-Based Graph Component
-    Pooling" <paper url>`_ paper.
-
+    Pooling" <https://arxiv.org/abs/2409.11856>`_ paper.
     :class:`ClusterPooling` computes a score for each edge.
     Based on the selected edges, graph clusters are calculated and compressed
     to one node using the injective :obj:`"sum"` aggregation function.
@@ -55,7 +54,7 @@ class ClusterPooling(torch.nn.Module):
         self.in_channels = in_channels
         self.edge_score_method = edge_score_method
         self.dropout = dropout
-        self.threshhold = threshold
+        self.threshold = threshold
 
         self.lin = torch.nn.Linear(2 * in_channels, 1)
 
@@ -116,7 +115,7 @@ class ClusterPooling(torch.nn.Module):
 
         from scipy.sparse.csgraph import connected_components
 
-        edge_contract = edge_index[:, edge_score > self.threshhold]
+        edge_contract = edge_index[:, edge_score > self.threshold]
 
         adj = to_scipy_sparse_matrix(edge_contract, num_nodes=x.size(0))
         _, cluster_np = connected_components(adj, directed=True,

torch_geometric/nn/pool/connect/base.py

@@ -4,7 +4,6 @@ from typing import Optional
 import torch
 from torch import Tensor
 
-import torch_geometric.typing
 from torch_geometric.nn.pool.select import SelectOutput
 
 
@@ -49,8 +48,7 @@ class ConnectOutput:
         self.batch = batch
 
 
-if torch_geometric.typing.WITH_PT113:
-    ConnectOutput = torch.jit.script(ConnectOutput)
+ConnectOutput = torch.jit.script(ConnectOutput)
 
 
 class Connect(torch.nn.Module):

torch_geometric/nn/pool/knn.py

@@ -91,9 +91,10 @@ class KNNIndex:
                     if hasattr(self.index, 'reserveMemory'):
                         self.index.reserveMemory(self.reserve)
                     else:
-                        warnings.warn(f"'{self.index.__class__.__name__}' "
-                                      f"does not support pre-allocation of "
-                                      f"memory")
+                        warnings.warn(
+                            f"'{self.index.__class__.__name__}' "
+                            f"does not support pre-allocation of "
+                            f"memory", stacklevel=2)
 
             self.index.train(emb)
 
@@ -135,14 +136,16 @@ class KNNIndex:
         query_k = min(query_k, self.numel)
 
         if k > 2048:  # `faiss` supports up-to `k=2048`:
-            warnings.warn(f"Capping 'k' to faiss' upper limit of 2048 "
-                          f"(got {k}). This may cause some relevant items to "
-                          f"not be retrieved.")
+            warnings.warn(
+                f"Capping 'k' to faiss' upper limit of 2048 "
+                f"(got {k}). This may cause some relevant items to "
+                f"not be retrieved.", stacklevel=2)
         elif query_k > 2048:
-            warnings.warn(f"Capping 'k' to faiss' upper limit of 2048 "
-                          f"(got {k} which got extended to {query_k} due to "
-                          f"the exclusion of existing links). This may cause "
-                          f"some relevant items to not be retrieved.")
+            warnings.warn(
+                f"Capping 'k' to faiss' upper limit of 2048 "
+                f"(got {k} which got extended to {query_k} due to "
+                f"the exclusion of existing links). This may cause "
+                f"some relevant items to not be retrieved.", stacklevel=2)
             query_k = 2048
 
         score, index = self.index.search(emb.detach(), query_k)

torch_geometric/nn/pool/select/base.py

@@ -4,8 +4,6 @@ from typing import Optional
 import torch
 from torch import Tensor
 
-import torch_geometric.typing
-
 
 @dataclass(init=False)
 class SelectOutput:
@@ -64,8 +62,7 @@ class SelectOutput:
         self.weight = weight
 
 
-if torch_geometric.typing.WITH_PT113:
-    SelectOutput = torch.jit.script(SelectOutput)
+SelectOutput = torch.jit.script(SelectOutput)
 
 
 class Select(torch.nn.Module):

torch_geometric/nn/to_hetero_module.py

@@ -108,9 +108,10 @@ class ToHeteroMessagePassing(torch.nn.Module):
 
         if (not hasattr(module, 'reset_parameters')
                 and sum([p.numel() for p in module.parameters()]) > 0):
-            warnings.warn(f"'{module}' will be duplicated, but its parameters "
-                          f"cannot be reset. To suppress this warning, add a "
-                          f"'reset_parameters()' method to '{module}'")
+            warnings.warn(
+                f"'{module}' will be duplicated, but its parameters "
+                f"cannot be reset. To suppress this warning, add a "
+                f"'reset_parameters()' method to '{module}'", stacklevel=2)
 
         convs = {edge_type: copy.deepcopy(module) for edge_type in edge_types}
         self.hetero_module = HeteroConv(convs, aggr)

torch_geometric/nn/to_hetero_transformer.py

@@ -157,7 +157,7 @@ class ToHeteroTransformer(Transformer):
                 f"There exist node types ({unused_node_types}) whose "
                 f"representations do not get updated during message passing "
                 f"as they do not occur as destination type in any edge type. "
-                f"This may lead to unexpected behavior.")
+                f"This may lead to unexpected behavior.", stacklevel=2)
 
         names = self.metadata[0] + [rel for _, rel, _ in self.metadata[1]]
         for name in names:
@@ -166,7 +166,7 @@ class ToHeteroTransformer(Transformer):
                     f"The type '{name}' contains invalid characters which "
                     f"may lead to unexpected behavior. To avoid any issues, "
                     f"ensure that your types only contain letters, numbers "
-                    f"and underscores.")
+                    f"and underscores.", stacklevel=2)
 
     def placeholder(self, node: Node, target: Any, name: str):
         # Adds a `get` call to the input dictionary for every node-type or
@@ -379,7 +379,7 @@ class ToHeteroTransformer(Transformer):
                 warnings.warn(
                     f"'{target}' will be duplicated, but its parameters "
                     f"cannot be reset. To suppress this warning, add a "
-                    f"'reset_parameters()' method to '{target}'")
+                    f"'reset_parameters()' method to '{target}'", stacklevel=2)
 
         return module_dict
 

torch_geometric/nn/to_hetero_with_bases_transformer.py

@@ -165,7 +165,7 @@ class ToHeteroWithBasesTransformer(Transformer):
                 f"There exist node types ({unused_node_types}) whose "
                 f"representations do not get updated during message passing "
                 f"as they do not occur as destination type in any edge type. "
-                f"This may lead to unexpected behavior.")
+                f"This may lead to unexpected behavior.", stacklevel=2)
 
         names = self.metadata[0] + [rel for _, rel, _ in self.metadata[1]]
         for name in names:
@@ -174,7 +174,7 @@ class ToHeteroWithBasesTransformer(Transformer):
                     f"The type '{name}' contains invalid characters which "
                     f"may lead to unexpected behavior. To avoid any issues, "
                     f"ensure that your types only contain letters, numbers "
-                    f"and underscores.")
+                    f"and underscores.", stacklevel=2)
 
     def transform(self) -> GraphModule:
         self._node_offset_dict_initialized = False
@@ -361,7 +361,7 @@ class HeteroBasisConv(torch.nn.Module):
                 warnings.warn(
                     f"'{conv}' will be duplicated, but its parameters cannot "
                     f"be reset. To suppress this warning, add a "
-                    f"'reset_parameters()' method to '{conv}'")
+                    f"'reset_parameters()' method to '{conv}'", stacklevel=2)
             torch.nn.init.xavier_uniform_(conv.edge_type_weight)
 
     def forward(self, edge_type: Tensor, *args, **kwargs) -> Tensor:
@@ -380,7 +380,7 @@ class HeteroBasisConv(torch.nn.Module):
 
 
 class LinearAlign(torch.nn.Module):
-    # Aligns representions to the same dimensionality. Note that this will
+    # Aligns representations to the same dimensionality. Note that this will
     # create lazy modules, and as such requires a forward pass in order to
     # initialize parameters.
     def __init__(self, keys: List[Union[NodeType, EdgeType]],

torch_geometric/profile/__init__.py

@@ -20,6 +20,7 @@ from .utils import (
     get_gpu_memory_from_nvidia_smi,
     get_model_size,
 )
+from .nvtx import nvtxit
 
 __all__ = [
     'profileit',
@@ -38,6 +39,7 @@ __all__ = [
     'get_gpu_memory_from_nvidia_smi',
     'get_gpu_memory_from_ipex',
     'benchmark',
+    'nvtxit',
 ]
 
 classes = __all__

torch_geometric/profile/nvtx.py

@@ -0,0 +1,66 @@
+from functools import wraps
+from typing import Optional
+
+import torch
+
+CUDA_PROFILE_STARTED = False
+
+
+def begin_cuda_profile():
+    global CUDA_PROFILE_STARTED
+    prev_state = CUDA_PROFILE_STARTED
+    if prev_state is False:
+        CUDA_PROFILE_STARTED = True
+        torch.cuda.cudart().cudaProfilerStart()
+    return prev_state
+
+
+def end_cuda_profile(prev_state: bool):
+    global CUDA_PROFILE_STARTED
+    CUDA_PROFILE_STARTED = prev_state
+    if prev_state is False:
+        torch.cuda.cudart().cudaProfilerStop()
+
+
+def nvtxit(name: Optional[str] = None, n_warmups: int = 0,
+           n_iters: Optional[int] = None):
+    """Enables NVTX profiling for a function.
+
+    Args:
+        name (Optional[str], optional): Name to give the reference frame for
+            the function being wrapped. Defaults to the name of the
+            function in code.
+        n_warmups (int, optional): Number of iters to call that function
+            before starting. Defaults to 0.
+        n_iters (Optional[int], optional): Number of iters of that function to
+            record. Defaults to all of them.
+    """
+    def nvtx(func):
+
+        nonlocal name
+        iters_so_far = 0
+        if name is None:
+            name = func.__name__
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            nonlocal iters_so_far
+            if not torch.cuda.is_available():
+                return func(*args, **kwargs)
+            elif iters_so_far < n_warmups:
+                iters_so_far += 1
+                return func(*args, **kwargs)
+            elif n_iters is None or iters_so_far < n_iters + n_warmups:
+                prev_state = begin_cuda_profile()
+                torch.cuda.nvtx.range_push(f"{name}_{iters_so_far}")
+                result = func(*args, **kwargs)
+                torch.cuda.nvtx.range_pop()
+                end_cuda_profile(prev_state)
+                iters_so_far += 1
+                return result
+            else:
+                return func(*args, **kwargs)
+
+        return wrapper
+
+    return nvtx

torch_geometric/profile/utils.py

@@ -119,19 +119,34 @@ def get_gpu_memory_from_nvidia_smi(  # pragma: no cover
         digits (int): The number of decimals to use for megabytes.
             (default: :obj:`2`)
     """
+    def parse_memory(output: str) -> list:
+        lines = output.decode('utf-8').split('\n')[1:-1]
+        mem_list = []
+        for line in lines:
+            try:
+                mem_list.append(int(line.split()[0]))
+            except (TypeError, ValueError):
+                mem_list.append(None)
+        return mem_list
+
+    def get_gpu_memory(out_device, digits):
+        if out_device is None:
+            return 0
+
+        return medibyte_to_megabyte(out_device, digits)
+
     CMD = 'nvidia-smi --query-gpu=memory.free --format=csv'
-    free_out = sp.check_output(CMD.split()).decode('utf-8').split('\n')[1:-1]
+    free_out = parse_memory(sp.check_output(CMD.split()))
 
     CMD = 'nvidia-smi --query-gpu=memory.used --format=csv'
-    used_out = sp.check_output(CMD.split()).decode('utf-8').split('\n')[1:-1]
+    used_out = parse_memory(sp.check_output(CMD.split()))
 
     if device < 0 or device >= len(free_out):
         raise AttributeError(
             f'GPU {device} not available (found {len(free_out)} GPUs)')
 
-    free_mem = medibyte_to_megabyte(int(free_out[device].split()[0]), digits)
-    used_mem = medibyte_to_megabyte(int(used_out[device].split()[0]), digits)
-
+    free_mem = get_gpu_memory(free_out[device], digits)
+    used_mem = get_gpu_memory(used_out[device], digits)
     return free_mem, used_mem
 
 

torch_geometric/sampler/__init__.py

@@ -3,7 +3,7 @@ r"""Graph sampler package."""
 from .base import (BaseSampler, NodeSamplerInput, EdgeSamplerInput,
                    SamplerOutput, HeteroSamplerOutput, NegativeSampling,
                    NumNeighbors)
-from .neighbor_sampler import NeighborSampler
+from .neighbor_sampler import NeighborSampler, BidirectionalNeighborSampler
 from .hgt_sampler import HGTSampler
 
 __all__ = classes = [
@@ -15,5 +15,6 @@ __all__ = classes = [
     'NumNeighbors',
     'NegativeSampling',
     'NeighborSampler',
+    'BidirectionalNeighborSampler',
     'HGTSampler',
 ]