pyg-nightly 2.7.0.dev20250701__tar.gz → 2.7.0.dev20250702__tar.gz

{pyg_nightly-2.7.0.dev20250701 → pyg_nightly-2.7.0.dev20250702}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyg-nightly
-Version: 2.7.0.dev20250701
+Version: 2.7.0.dev20250702
 Summary: Graph Neural Network Library for PyTorch
 Keywords: deep-learning,pytorch,geometric-deep-learning,graph-neural-networks,graph-convolutional-networks
 Author-email: Matthias Fey <matthias@pyg.org>

{pyg_nightly-2.7.0.dev20250701 → pyg_nightly-2.7.0.dev20250702}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend="flit_core.buildapi"
 
 [project]
 name="pyg-nightly"
-version="2.7.0.dev20250701"
+version="2.7.0.dev20250702"
 authors=[
     {name="Matthias Fey", email="matthias@pyg.org"},
 ]

{pyg_nightly-2.7.0.dev20250701 → pyg_nightly-2.7.0.dev20250702}/torch_geometric/__init__.py

@@ -31,7 +31,7 @@ from .lazy_loader import LazyLoader
 contrib = LazyLoader('contrib', globals(), 'torch_geometric.contrib')
 graphgym = LazyLoader('graphgym', globals(), 'torch_geometric.graphgym')
 
-__version__ = '2.7.0.dev20250701'
+__version__ = '2.7.0.dev20250702'
 
 __all__ = [
     'Index',

{pyg_nightly-2.7.0.dev20250701 → pyg_nightly-2.7.0.dev20250702}/torch_geometric/nn/aggr/patch_transformer.py

@@ -32,6 +32,8 @@ class PatchTransformerAggregation(Aggregation):
         aggr (str or list[str], optional): The aggregation module, *e.g.*,
             :obj:`"sum"`, :obj:`"mean"`, :obj:`"min"`, :obj:`"max"`,
             :obj:`"var"`, :obj:`"std"`. (default: :obj:`"mean"`)
+        device (torch.device, optional): The device of the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -43,6 +45,7 @@ class PatchTransformerAggregation(Aggregation):
         heads: int = 1,
         dropout: float = 0.0,
         aggr: Union[str, List[str]] = 'mean',
+        device: Optional[torch.device] = None,
     ) -> None:
         super().__init__()
 
@@ -55,12 +58,13 @@ class PatchTransformerAggregation(Aggregation):
         for aggr in self.aggrs:
             assert aggr in ['sum', 'mean', 'min', 'max', 'var', 'std']
 
-        self.lin = torch.nn.Linear(in_channels, hidden_channels)
+        self.lin = torch.nn.Linear(in_channels, hidden_channels, device=device)
         self.pad_projector = torch.nn.Linear(
             patch_size * hidden_channels,
             hidden_channels,
+            device=device,
         )
-        self.pe = PositionalEncoding(hidden_channels)
+        self.pe = PositionalEncoding(hidden_channels, device=device)
 
         self.blocks = torch.nn.ModuleList([
             MultiheadAttentionBlock(
@@ -68,12 +72,14 @@ class PatchTransformerAggregation(Aggregation):
                 heads=heads,
                 layer_norm=True,
                 dropout=dropout,
+                device=device,
             ) for _ in range(num_transformer_blocks)
         ])
 
         self.fc = torch.nn.Linear(
             hidden_channels * len(self.aggrs),
             out_channels,
+            device=device,
         )
 
     def reset_parameters(self) -> None:

{pyg_nightly-2.7.0.dev20250701 → pyg_nightly-2.7.0.dev20250702}/torch_geometric/nn/aggr/utils.py

@@ -26,9 +26,11 @@ class MultiheadAttentionBlock(torch.nn.Module):
             normalization. (default: :obj:`True`)
         dropout (float, optional): Dropout probability of attention weights.
             (default: :obj:`0`)
+        device (torch.device, optional): The device of the module.
+            (default: :obj:`None`)
     """
     def __init__(self, channels: int, heads: int = 1, layer_norm: bool = True,
-                 dropout: float = 0.0):
+                 dropout: float = 0.0, device: Optional[torch.device] = None):
         super().__init__()
 
         self.channels = channels
@@ -40,10 +42,13 @@ class MultiheadAttentionBlock(torch.nn.Module):
             heads,
             batch_first=True,
             dropout=dropout,
+            device=device,
         )
-        self.lin = Linear(channels, channels)
-        self.layer_norm1 = LayerNorm(channels) if layer_norm else None
-        self.layer_norm2 = LayerNorm(channels) if layer_norm else None
+        self.lin = Linear(channels, channels, device=device)
+        self.layer_norm1 = LayerNorm(channels,
+                                     device=device) if layer_norm else None
+        self.layer_norm2 = LayerNorm(channels,
+                                     device=device) if layer_norm else None
 
     def reset_parameters(self):
         self.attn._reset_parameters()

{pyg_nightly-2.7.0.dev20250701 → pyg_nightly-2.7.0.dev20250702}/torch_geometric/nn/encoding.py

@@ -1,4 +1,5 @@
 import math
+from typing import Optional
 
 import torch
 from torch import Tensor
@@ -23,12 +24,15 @@ class PositionalEncoding(torch.nn.Module):
         granularity (float, optional): The granularity of the positions. If
             set to smaller value, the encoder will capture more fine-grained
             changes in positions. (default: :obj:`1.0`)
+        device (torch.device, optional): The device of the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
         out_channels: int,
         base_freq: float = 1e-4,
         granularity: float = 1.0,
+        device: Optional[torch.device] = None,
     ):
         super().__init__()
 
@@ -40,7 +44,8 @@ class PositionalEncoding(torch.nn.Module):
         self.base_freq = base_freq
         self.granularity = granularity
 
-        frequency = torch.logspace(0, 1, out_channels // 2, base_freq)
+        frequency = torch.logspace(0, 1, out_channels // 2, base_freq,
+                                   device=device)
         self.register_buffer('frequency', frequency)
 
         self.reset_parameters()
@@ -75,13 +80,17 @@ class TemporalEncoding(torch.nn.Module):
 
     Args:
         out_channels (int): Size :math:`d` of each output sample.
+        device (torch.device, optional): The device of the module.
+            (default: :obj:`None`)
     """
-    def __init__(self, out_channels: int):
+    def __init__(self, out_channels: int,
+                 device: Optional[torch.device] = None):
         super().__init__()
         self.out_channels = out_channels
 
         sqrt = math.sqrt(out_channels)
-        weight = 1.0 / sqrt**torch.linspace(0, sqrt, out_channels).view(1, -1)
+        weight = 1.0 / sqrt**torch.linspace(0, sqrt, out_channels,
+                                            device=device).view(1, -1)
         self.register_buffer('weight', weight)
 
         self.reset_parameters()

{pyg_nightly-2.7.0.dev20250701 → pyg_nightly-2.7.0.dev20250702}/torch_geometric/nn/norm/batch_norm.py

@@ -39,6 +39,8 @@ class BatchNorm(torch.nn.Module):
             with only a single element will work as during in evaluation.
             That is the running mean and variance will be used.
             Requires :obj:`track_running_stats=True`. (default: :obj:`False`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -48,6 +50,7 @@ class BatchNorm(torch.nn.Module):
         affine: bool = True,
         track_running_stats: bool = True,
         allow_single_element: bool = False,
+        device: Optional[torch.device] = None,
     ):
         super().__init__()
 
@@ -56,7 +59,7 @@ class BatchNorm(torch.nn.Module):
                              "'track_running_stats' to be set to `True`")
 
         self.module = torch.nn.BatchNorm1d(in_channels, eps, momentum, affine,
-                                           track_running_stats)
+                                           track_running_stats, device=device)
         self.in_channels = in_channels
         self.allow_single_element = allow_single_element
 
@@ -114,6 +117,8 @@ class HeteroBatchNorm(torch.nn.Module):
             :obj:`False`, this module does not track such statistics and always
             uses batch statistics in both training and eval modes.
             (default: :obj:`True`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -123,6 +128,7 @@ class HeteroBatchNorm(torch.nn.Module):
         momentum: Optional[float] = 0.1,
         affine: bool = True,
         track_running_stats: bool = True,
+        device: Optional[torch.device] = None,
     ):
         super().__init__()
 
@@ -134,17 +140,21 @@ class HeteroBatchNorm(torch.nn.Module):
         self.track_running_stats = track_running_stats
 
         if self.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)
 
         if self.track_running_stats:
-            self.register_buffer('running_mean',
-                                 torch.empty(num_types, in_channels))
-            self.register_buffer('running_var',
-                                 torch.empty(num_types, in_channels))
+            self.register_buffer(
+                'running_mean',
+                torch.empty(num_types, in_channels, device=device))
+            self.register_buffer(
+                'running_var',
+                torch.empty(num_types, in_channels, device=device))
             self.register_buffer('num_batches_tracked', torch.tensor(0))
         else:
             self.register_buffer('running_mean', None)

{pyg_nightly-2.7.0.dev20250701 → pyg_nightly-2.7.0.dev20250702}/torch_geometric/nn/norm/diff_group_norm.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import torch
 from torch import Tensor
 from torch.nn import BatchNorm1d, Linear
@@ -39,6 +41,8 @@ class DiffGroupNorm(torch.nn.Module):
             :obj:`False`, this module does not track such statistics and always
             uses batch statistics in both training and eval modes.
             (default: :obj:`True`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -49,6 +53,7 @@ class DiffGroupNorm(torch.nn.Module):
         momentum: float = 0.1,
         affine: bool = True,
         track_running_stats: bool = True,
+        device: Optional[torch.device] = None,
     ):
         super().__init__()
 
@@ -56,9 +61,9 @@ class DiffGroupNorm(torch.nn.Module):
         self.groups = groups
         self.lamda = lamda
 
-        self.lin = Linear(in_channels, groups, bias=False)
+        self.lin = Linear(in_channels, groups, bias=False, device=device)
         self.norm = BatchNorm1d(groups * in_channels, eps, momentum, affine,
-                                track_running_stats)
+                                track_running_stats, device=device)
 
         self.reset_parameters()
 

{pyg_nightly-2.7.0.dev20250701 → pyg_nightly-2.7.0.dev20250702}/torch_geometric/nn/norm/graph_norm.py

@@ -26,16 +26,21 @@ class GraphNorm(torch.nn.Module):
         in_channels (int): Size of each input sample.
         eps (float, optional): A value added to the denominator for numerical
             stability. (default: :obj:`1e-5`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
-    def __init__(self, in_channels: int, eps: float = 1e-5):
+    def __init__(self, in_channels: int, eps: float = 1e-5,
+                 device: Optional[torch.device] = None):
         super().__init__()
 
         self.in_channels = in_channels
         self.eps = eps
 
-        self.weight = torch.nn.Parameter(torch.empty(in_channels))
-        self.bias = torch.nn.Parameter(torch.empty(in_channels))
-        self.mean_scale = torch.nn.Parameter(torch.empty(in_channels))
+        self.weight = torch.nn.Parameter(
+            torch.empty(in_channels, device=device))
+        self.bias = torch.nn.Parameter(torch.empty(in_channels, device=device))
+        self.mean_scale = torch.nn.Parameter(
+            torch.empty(in_channels, device=device))
 
         self.reset_parameters()
 

{pyg_nightly-2.7.0.dev20250701 → pyg_nightly-2.7.0.dev20250702}/torch_geometric/nn/norm/instance_norm.py

@@ -1,5 +1,6 @@
 from typing import Optional
 
+import torch
 import torch.nn.functional as F
 from torch import Tensor
 from torch.nn.modules.instancenorm import _InstanceNorm
@@ -36,6 +37,8 @@ class InstanceNorm(_InstanceNorm):
             :obj:`False`, this module does not track such statistics and always
             uses instance statistics in both training and eval modes.
             (default: :obj:`False`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -44,9 +47,10 @@ class InstanceNorm(_InstanceNorm):
         momentum: float = 0.1,
         affine: bool = False,
         track_running_stats: bool = False,
+        device: Optional[torch.device] = None,
     ):
         super().__init__(in_channels, eps, momentum, affine,
-                         track_running_stats)
+                         track_running_stats, device=device)
 
     def reset_parameters(self):
         r"""Resets all learnable parameters of the module."""

{pyg_nightly-2.7.0.dev20250701 → pyg_nightly-2.7.0.dev20250702}/torch_geometric/nn/norm/layer_norm.py

@@ -35,6 +35,8 @@ class LayerNorm(torch.nn.Module):
             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)
@@ -134,6 +137,8 @@ class HeteroLayerNorm(torch.nn.Module):
             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)

{pyg_nightly-2.7.0.dev20250701 → pyg_nightly-2.7.0.dev20250702}/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):