pyg-nightly 2.6.0.dev20240704__py3-none-any.whl → 2.8.0.dev20251207__py3-none-any.whl

torch_geometric/nn/encoding.py

@@ -1,8 +1,14 @@
 import math
+from typing import Optional
 
 import torch
 from torch import Tensor
 
+__all__ = classes = [
+    'PositionalEncoding',
+    'TemporalEncoding',
+]
+
 
 class PositionalEncoding(torch.nn.Module):
     r"""The positional encoding scheme from the `"Attention Is All You Need"
@@ -23,12 +29,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 +49,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 +85,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()

torch_geometric/nn/fx.py

@@ -1,8 +1,9 @@
 import copy
 import warnings
-from typing import Any, Dict, Optional
+from typing import Any, Callable, Dict, List, Optional, Type, Union
 
 import torch
+from torch import Tensor
 from torch.nn import Module, ModuleDict, ModuleList, Sequential
 
 try:
@@ -18,8 +19,8 @@ class Transformer:
     :class:`~torch.nn.Module`.
     :class:`Transformer` works entirely symbolically.
 
-    Methods in the :class:`Transformer` class can be overriden to customize the
-    behavior of transformation.
+    Methods in the :class:`Transformer` class can be overridden to customize
+    the behavior of transformation.
 
     .. code-block:: none
 
@@ -129,11 +130,13 @@ class Transformer:
         # (node-level, edge-level) by filling `self._state`:
         for node in list(self.graph.nodes):
             if node.op == 'call_function' and 'training' in node.kwargs:
-                warnings.warn(f"Found function '{node.name}' with keyword "
-                              f"argument 'training'. During FX tracing, this "
-                              f"will likely be baked in as a constant value. "
-                              f"Consider replacing this function by a module "
-                              f"to properly encapsulate its training flag.")
+                warnings.warn(
+                    f"Found function '{node.name}' with keyword "
+                    f"argument 'training'. During FX tracing, this "
+                    f"will likely be baked in as a constant value. "
+                    f"Consider replacing this function by a module "
+                    f"to properly encapsulate its training flag.",
+                    stacklevel=2)
 
             if node.op == 'placeholder':
                 if node.name not in self._state:
@@ -283,13 +286,13 @@ def symbolic_trace(
             # TODO We currently only trace top-level modules.
             return not isinstance(module, torch.nn.Sequential)
 
-        # Note: This is a hack around the fact that `Aggregaton.__call__`
+        # Note: This is a hack around the fact that `Aggregation.__call__`
         # is not patched by the base implementation of `trace`.
         # see https://github.com/pyg-team/pytorch_geometric/pull/5021 for
         # details on the rationale
         # TODO: Revisit https://github.com/pyg-team/pytorch_geometric/pull/5021
         @st.compatibility(is_backward_compatible=True)
-        def trace(self, root: st.Union[torch.nn.Module, st.Callable[..., Any]],
+        def trace(self, root: Union[torch.nn.Module, Callable[..., Any]],
                   concrete_args: Optional[Dict[str, Any]] = None) -> Graph:
 
             if isinstance(root, torch.nn.Module):
@@ -303,17 +306,16 @@ def symbolic_trace(
                 self.root = torch.nn.Module()
                 fn = root
 
-            tracer_cls: Optional[st.Type['Tracer']] = getattr(
+            tracer_cls: Optional[Type['Tracer']] = getattr(
                 self, '__class__', None)
             self.graph = Graph(tracer_cls=tracer_cls)
 
-            self.tensor_attrs: Dict[st.Union[torch.Tensor, st.ScriptObject],
-                                    str] = {}
+            self.tensor_attrs: Dict[Union[Tensor, st.ScriptObject], str] = {}
 
             def collect_tensor_attrs(m: torch.nn.Module,
-                                     prefix_atoms: st.List[str]):
+                                     prefix_atoms: List[str]):
                 for k, v in m.__dict__.items():
-                    if isinstance(v, (torch.Tensor, st.ScriptObject)):
+                    if isinstance(v, (Tensor, st.ScriptObject)):
                         self.tensor_attrs[v] = '.'.join(prefix_atoms + [k])
                 for k, v in m.named_children():
                     collect_tensor_attrs(v, prefix_atoms + [k])

torch_geometric/nn/model_hub.py

@@ -4,6 +4,8 @@ from typing import Any, Dict, Optional, Union
 
 import torch
 
+from torch_geometric.io import fs
+
 try:
     from huggingface_hub import ModelHubMixin, hf_hub_download
 except ImportError:
@@ -142,10 +144,10 @@ class PyGModelHubMixin(ModelHubMixin):
         revision,
         cache_dir,
         force_download,
-        proxies,
-        resume_download,
         local_files_only,
         token,
+        proxies=None,
+        resume_download=False,
         dataset_name='',
         model_name='',
         map_location='cpu',
@@ -163,8 +165,6 @@ class PyGModelHubMixin(ModelHubMixin):
                 revision=revision,
                 cache_dir=cache_dir,
                 force_download=force_download,
-                proxies=proxies,
-                resume_download=resume_download,
                 token=token,
                 local_files_only=local_files_only,
             )
@@ -175,7 +175,7 @@ class PyGModelHubMixin(ModelHubMixin):
 
         model = cls(dataset_name, model_name, model_kwargs)
 
-        state_dict = torch.load(model_file, map_location=map_location)
+        state_dict = fs.torch_load(model_file, map_location=map_location)
         model.load_state_dict(state_dict, strict=strict)
         model.eval()
 
@@ -186,8 +186,6 @@ class PyGModelHubMixin(ModelHubMixin):
         cls,
         pretrained_model_name_or_path: str,
         force_download: bool = False,
-        resume_download: bool = False,
-        proxies: Optional[Dict] = None,
         token: Optional[Union[str, bool]] = None,
         cache_dir: Optional[str] = None,
         local_files_only: bool = False,
@@ -213,13 +211,6 @@ class PyGModelHubMixin(ModelHubMixin):
                 (re-)download of the model weights and configuration files,
                 overriding the cached versions if they exist.
                 (default: :obj:`False`)
-            resume_download (bool, optional): Whether to delete incompletely
-                received files. Will attempt to resume the download if such a
-                file exists. (default: :obj:`False`)
-            proxies (Dict[str, str], optional): A dictionary of proxy servers
-                to use by protocol or endpoint, *e.g.*,
-                :obj:`{'http': 'foo.bar:3128', 'http://host': 'foo.bar:4012'}`.
-                The proxies are used on each request. (default: :obj:`None`)
             token (str or bool, optional): The token to use as HTTP bearer
                 authorization for remote files. If set to :obj:`True`, will use
                 the token generated when running :obj:`transformers-cli login`
@@ -237,8 +228,6 @@ class PyGModelHubMixin(ModelHubMixin):
         return super().from_pretrained(
             pretrained_model_name_or_path,
             force_download=force_download,
-            resume_download=resume_download,
-            proxies=proxies,
             use_auth_token=token,
             cache_dir=cache_dir,
             local_files_only=local_files_only,

torch_geometric/nn/models/__init__.py

@@ -12,6 +12,7 @@ from .re_net import RENet
 from .graph_unet import GraphUNet
 from .schnet import SchNet
 from .dimenet import DimeNet, DimeNetPlusPlus
+from .gpse import GPSE, GPSENodeEncoder
 from .captum import to_captum_model
 from .metapath2vec import MetaPath2Vec
 from .deepgcn import DeepGCNLayer
@@ -28,10 +29,14 @@ from .gnnff import GNNFF
 from .pmlp import PMLP
 from .neural_fingerprint import NeuralFingerprint
 from .visnet import ViSNet
+from .lpformer import LPFormer
+from .sgformer import SGFormer
 
+from .polynormer import Polynormer
 # Deprecated:
 from torch_geometric.explain.algorithm.captum import (to_captum_input,
                                                       captum_output_to_dicts)
+from .attract_repel import ARLinkPredictor
 
 __all__ = classes = [
     'MLP',
@@ -57,6 +62,8 @@ __all__ = classes = [
     'SchNet',
     'DimeNet',
     'DimeNetPlusPlus',
+    'GPSE',
+    'GPSENodeEncoder',
     'to_captum_model',
     'to_captum_input',
     'captum_output_to_dicts',
@@ -75,4 +82,8 @@ __all__ = classes = [
     'PMLP',
     'NeuralFingerprint',
     'ViSNet',
+    'LPFormer',
+    'SGFormer',
+    'Polynormer',
+    'ARLinkPredictor',
 ]

torch_geometric/nn/models/attentive_fp.py

@@ -160,7 +160,7 @@ class AttentiveFP(torch.nn.Module):
         edge_index = torch.stack([row, batch], dim=0)
 
         out = global_add_pool(x, batch).relu_()
-        for t in range(self.num_timesteps):
+        for _ in range(self.num_timesteps):
             h = F.elu_(self.mol_conv((x, out), edge_index))
             h = F.dropout(h, p=self.dropout, training=self.training)
             out = self.mol_gru(h, out).relu_()

torch_geometric/nn/models/attract_repel.py

@@ -0,0 +1,148 @@
+import torch
+import torch.nn.functional as F
+
+
+class ARLinkPredictor(torch.nn.Module):
+    r"""Link predictor using Attract-Repel embeddings from the paper
+    `"Pseudo-Euclidean Attract-Repel Embeddings for Undirected Graphs"
+    <https://arxiv.org/abs/2106.09671>`_.
+
+    This model splits node embeddings into: attract and
+    repel.
+    The edge prediction score is computed as the dot product of attract
+    components minus the dot product of repel components.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        hidden_channels (int): Size of hidden embeddings.
+        out_channels (int, optional): Size of output embeddings.
+            If set to :obj:`None`, will default to :obj:`hidden_channels`.
+            (default: :obj:`None`)
+        num_layers (int): Number of message passing layers.
+            (default: :obj:`2`)
+        dropout (float): Dropout probability. (default: :obj:`0.0`)
+        attract_ratio (float): Ratio to use for attract component.
+            Must be between 0 and 1. (default: :obj:`0.5`)
+    """
+    def __init__(self, in_channels, hidden_channels, out_channels=None,
+                 num_layers=2, dropout=0.0, attract_ratio=0.5):
+        super().__init__()
+
+        if out_channels is None:
+            out_channels = hidden_channels
+
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.num_layers = num_layers
+        self.dropout = dropout
+
+        if not 0 <= attract_ratio <= 1:
+            raise ValueError(
+                f"attract_ratio must be between 0 and 1, got {attract_ratio}")
+
+        self.attract_ratio = attract_ratio
+        self.attract_dim = int(out_channels * attract_ratio)
+        self.repel_dim = out_channels - self.attract_dim
+
+        # Create model layers
+        self.lins = torch.nn.ModuleList()
+        self.lins.append(torch.nn.Linear(in_channels, hidden_channels))
+
+        for _ in range(num_layers - 2):
+            self.lins.append(torch.nn.Linear(hidden_channels, hidden_channels))
+
+        # Final layer splits into attract and repel components
+        self.lin_attract = torch.nn.Linear(hidden_channels, self.attract_dim)
+        self.lin_repel = torch.nn.Linear(hidden_channels, self.repel_dim)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        """Reset all learnable parameters."""
+        for lin in self.lins:
+            lin.reset_parameters()
+        self.lin_attract.reset_parameters()
+        self.lin_repel.reset_parameters()
+
+    def encode(self, x, *args, **kwargs):
+        """Encode node features into attract-repel embeddings.
+
+        Args:
+            x (torch.Tensor): Node feature matrix of shape
+                :obj:`[num_nodes, in_channels]`.
+            *args: Variable length argument list
+            **kwargs: Arbitrary keyword arguments
+
+        """
+        for lin in self.lins:
+            x = lin(x)
+            x = F.relu(x)
+            x = F.dropout(x, p=self.dropout, training=self.training)
+
+        # Split into attract and repel components
+        attract_x = self.lin_attract(x)
+        repel_x = self.lin_repel(x)
+
+        return attract_x, repel_x
+
+    def decode(self, attract_z, repel_z, edge_index):
+        """Decode edge scores from attract-repel embeddings.
+
+        Args:
+            attract_z (torch.Tensor): Attract embeddings of shape
+                :obj:`[num_nodes, attract_dim]`.
+            repel_z (torch.Tensor): Repel embeddings of shape
+                :obj:`[num_nodes, repel_dim]`.
+            edge_index (torch.Tensor): Edge indices of shape
+                :obj:`[2, num_edges]`.
+
+        Returns:
+            torch.Tensor: Edge prediction scores.
+        """
+        # Get node embeddings for edges
+        row, col = edge_index
+        attract_z_row = attract_z[row]
+        attract_z_col = attract_z[col]
+        repel_z_row = repel_z[row]
+        repel_z_col = repel_z[col]
+
+        # Compute attract-repel scores
+        attract_score = torch.sum(attract_z_row * attract_z_col, dim=1)
+        repel_score = torch.sum(repel_z_row * repel_z_col, dim=1)
+
+        return attract_score - repel_score
+
+    def forward(self, x, edge_index):
+        """Forward pass for link prediction.
+
+        Args:
+            x (torch.Tensor): Node feature matrix.
+            edge_index (torch.Tensor): Edge indices to predict.
+
+        Returns:
+            torch.Tensor: Predicted edge scores.
+        """
+        # Encode nodes into attract-repel embeddings
+        attract_z, repel_z = self.encode(x)
+
+        # Decode target edges
+        return torch.sigmoid(self.decode(attract_z, repel_z, edge_index))
+
+    def calculate_r_fraction(self, attract_z, repel_z):
+        """Calculate the R-fraction (proportion of energy in repel space).
+
+        Args:
+            attract_z (torch.Tensor): Attract embeddings.
+            repel_z (torch.Tensor): Repel embeddings.
+
+        Returns:
+            float: R-fraction value.
+        """
+        attract_norm_squared = torch.sum(attract_z**2)
+        repel_norm_squared = torch.sum(repel_z**2)
+
+        r_fraction = repel_norm_squared / (attract_norm_squared +
+                                           repel_norm_squared + 1e-10)
+
+        return r_fraction.item()

torch_geometric/nn/models/basic_gnn.py

@@ -415,7 +415,8 @@ class GCN(BasicGNN):
             (default: :obj:`None`)
         jk (str, optional): The Jumping Knowledge mode. If specified, the model
             will additionally apply a final linear transformation to transform
-            node embeddings to the expected output feature dimensionality.
+            node embeddings to the expected output feature dimensionality,
+            while default will not.
             (:obj:`None`, :obj:`"last"`, :obj:`"cat"`, :obj:`"max"`,
             :obj:`"lstm"`). (default: :obj:`None`)
         **kwargs (optional): Additional arguments of

torch_geometric/nn/models/captum.py

@@ -94,7 +94,7 @@ def to_captum_model(
             function will return the output of the model for the element at
             the index specified. (default: :obj:`None`)
         metadata (Metadata, optional): The metadata of the heterogeneous graph.
-            Only required if explaning a
+            Only required if explaining a
             :class:`~torch_geometric.data.HeteroData` object.
             (default: :obj:`None`)
     """

torch_geometric/nn/models/deep_graph_infomax.py

@@ -106,7 +106,7 @@ class DeepGraphInfomax(torch.nn.Module):
         """
         from sklearn.linear_model import LogisticRegression
 
-        clf = LogisticRegression(solver=solver, *args,
+        clf = LogisticRegression(*args, solver=solver,
                                  **kwargs).fit(train_z.detach().cpu().numpy(),
                                                train_y.detach().cpu().numpy())
         return clf.score(test_z.detach().cpu().numpy(),

torch_geometric/nn/models/dimenet.py

@@ -755,7 +755,7 @@ class DimeNetPlusPlus(DimeNet):
             interaction blocks after the skip connection. (default: :obj:`2`)
         num_output_layers: (int, optional): Number of linear layers for the
             output blocks. (default: :obj:`3`)
-        act: (str or Callable, optional): The activation funtion.
+        act: (str or Callable, optional): The activation function.
             (default: :obj:`"swish"`)
         output_initializer (str, optional): The initialization method for the
             output layer (:obj:`"zeros"`, :obj:`"glorot_orthogonal"`).
@@ -805,7 +805,7 @@ class DimeNetPlusPlus(DimeNet):
 
         # We are re-using the RBF, SBF and embedding layers of `DimeNet` and
         # redefine output_block and interaction_block in DimeNet++.
-        # Hence, it is to be noted that in the above initalization, the
+        # Hence, it is to be noted that in the above initialization, the
         # variable `num_bilinear` does not have any purpose as it is used
         # solely in the `OutputBlock` of DimeNet:
         self.output_blocks = torch.nn.ModuleList([

torch_geometric/nn/models/dimenet_utils.py

@@ -1,5 +1,7 @@
 # Shameless steal from: https://github.com/klicperajo/dimenet
 
+import math
+
 import numpy as np
 import sympy as sym
 from scipy import special as sp
@@ -62,8 +64,8 @@ def bessel_basis(n, k):
 
 
 def sph_harm_prefactor(k, m):
-    return ((2 * k + 1) * np.math.factorial(k - abs(m)) /
-            (4 * np.pi * np.math.factorial(k + abs(m))))**0.5
+    return ((2 * k + 1) * math.factorial(k - abs(m)) /
+            (4 * np.pi * math.factorial(k + abs(m))))**0.5
 
 
 def associated_legendre_polynomials(k, zero_m_only=True):