pyg-nightly 2.6.0.dev20240511__py3-none-any.whl → 2.7.0.dev20250114__py3-none-any.whl

torch_geometric/nn/conv/message_passing.py

@@ -167,71 +167,8 @@ class MessagePassing(torch.nn.Module):
         self._edge_update_forward_pre_hooks: HookDict = OrderedDict()
         self._edge_update_forward_hooks: HookDict = OrderedDict()
 
-        root_dir = osp.dirname(osp.realpath(__file__))
-        jinja_prefix = f'{self.__module__}_{self.__class__.__name__}'
-        # Optimize `propagate()` via `*.jinja` templates:
-        if not self.propagate.__module__.startswith(jinja_prefix):
-            try:
-                if 'propagate' in self.__class__.__dict__:
-                    raise ValueError("Cannot compile custom 'propagate' "
-                                     "method")
-
-                module = module_from_template(
-                    module_name=f'{jinja_prefix}_propagate',
-                    template_path=osp.join(root_dir, 'propagate.jinja'),
-                    tmp_dirname='message_passing',
-                    # Keyword arguments:
-                    modules=self.inspector._modules,
-                    collect_name='collect',
-                    signature=self._get_propagate_signature(),
-                    collect_param_dict=self.inspector.get_flat_param_dict(
-                        ['message', 'aggregate', 'update']),
-                    message_args=self.inspector.get_param_names('message'),
-                    aggregate_args=self.inspector.get_param_names('aggregate'),
-                    message_and_aggregate_args=self.inspector.get_param_names(
-                        'message_and_aggregate'),
-                    update_args=self.inspector.get_param_names('update'),
-                    fuse=self.fuse,
-                )
-
-                self.__class__._orig_propagate = self.__class__.propagate
-                self.__class__._jinja_propagate = module.propagate
-
-                self.__class__.propagate = module.propagate
-                self.__class__.collect = module.collect
-            except Exception:  # pragma: no cover
-                self.__class__._orig_propagate = self.__class__.propagate
-                self.__class__._jinja_propagate = self.__class__.propagate
-
-        # Optimize `edge_updater()` via `*.jinja` templates (if implemented):
-        if (self.inspector.implements('edge_update')
-                and not self.edge_updater.__module__.startswith(jinja_prefix)):
-            try:
-                if 'edge_updater' in self.__class__.__dict__:
-                    raise ValueError("Cannot compile custom 'edge_updater' "
-                                     "method")
-
-                module = module_from_template(
-                    module_name=f'{jinja_prefix}_edge_updater',
-                    template_path=osp.join(root_dir, 'edge_updater.jinja'),
-                    tmp_dirname='message_passing',
-                    # Keyword arguments:
-                    modules=self.inspector._modules,
-                    collect_name='edge_collect',
-                    signature=self._get_edge_updater_signature(),
-                    collect_param_dict=self.inspector.get_param_dict(
-                        'edge_update'),
-                )
-
-                self.__class__._orig_edge_updater = self.__class__.edge_updater
-                self.__class__._jinja_edge_updater = module.edge_updater
-
-                self.__class__.edge_updater = module.edge_updater
-                self.__class__.edge_collect = module.edge_collect
-            except Exception:  # pragma: no cover
-                self.__class__._orig_edge_updater = self.__class__.edge_updater
-                self.__class__._jinja_edge_updater = (
-                    self.__class__.edge_updater)
+        # Set jittable `propagate` and `edge_updater` function templates:
+        self._set_jittable_templates()
 
         # Explainability:
         self._explain: Optional[bool] = None
@@ -248,6 +185,12 @@ class MessagePassing(torch.nn.Module):
         if self.aggr_module is not None:
             self.aggr_module.reset_parameters()
 
+    def __setstate__(self, data: Dict[str, Any]) -> None:
+        self.inspector = data['inspector']
+        self.fuse = data['fuse']
+        self._set_jittable_templates()
+        super().__setstate__(data)
+
     def __repr__(self) -> str:
         channels_repr = ''
         if hasattr(self, 'in_channels') and hasattr(self, 'out_channels'):
@@ -261,7 +204,7 @@ class MessagePassing(torch.nn.Module):
     def _check_input(
         self,
         edge_index: Union[Tensor, SparseTensor],
-        size: Optional[Tuple[int, int]],
+        size: Optional[Tuple[Optional[int], Optional[int]]],
     ) -> List[Optional[int]]:
 
         if not torch.jit.is_scripting() and isinstance(edge_index, EdgeIndex):
@@ -270,19 +213,20 @@ class MessagePassing(torch.nn.Module):
         if is_sparse(edge_index):
             if self.flow == 'target_to_source':
                 raise ValueError(
-                    ('Flow direction "target_to_source" is invalid for '
-                     'message propagation via `torch_sparse.SparseTensor` '
-                     'or `torch.sparse.Tensor`. If you really want to make '
-                     'use of a reverse message passing flow, pass in the '
-                     'transposed sparse tensor to the message passing module, '
-                     'e.g., `adj_t.t()`.'))
+                    'Flow direction "target_to_source" is invalid for '
+                    'message propagation via `torch_sparse.SparseTensor` '
+                    'or `torch.sparse.Tensor`. If you really want to make '
+                    'use of a reverse message passing flow, pass in the '
+                    'transposed sparse tensor to the message passing module, '
+                    'e.g., `adj_t.t()`.')
 
             if isinstance(edge_index, SparseTensor):
                 return [edge_index.size(1), edge_index.size(0)]
             return [edge_index.size(1), edge_index.size(0)]
 
         elif isinstance(edge_index, Tensor):
-            int_dtypes = (torch.uint8, torch.int8, torch.int32, torch.int64)
+            int_dtypes = (torch.uint8, torch.int8, torch.int16, torch.int32,
+                          torch.int64)
 
             if edge_index.dtype not in int_dtypes:
                 raise ValueError(f"Expected 'edge_index' to be of integer "
@@ -298,9 +242,9 @@ class MessagePassing(torch.nn.Module):
             return list(size) if size is not None else [None, None]
 
         raise ValueError(
-            ('`MessagePassing.propagate` only supports integer tensors of '
-             'shape `[2, num_messages]`, `torch_sparse.SparseTensor` or '
-             '`torch.sparse.Tensor` for argument `edge_index`.'))
+            '`MessagePassing.propagate` only supports integer tensors of '
+            'shape `[2, num_messages]`, `torch_sparse.SparseTensor` or '
+            '`torch.sparse.Tensor` for argument `edge_index`.')
 
     def _set_size(
         self,
@@ -313,8 +257,8 @@ class MessagePassing(torch.nn.Module):
             size[dim] = src.size(self.node_dim)
         elif the_size != src.size(self.node_dim):
             raise ValueError(
-                (f'Encountered tensor with size {src.size(self.node_dim)} in '
-                 f'dimension {self.node_dim}, but expected size {the_size}.'))
+                f'Encountered tensor with size {src.size(self.node_dim)} in '
+                f'dimension {self.node_dim}, but expected size {the_size}.')
 
     def _index_select(self, src: Tensor, index) -> Tensor:
         if torch.jit.is_scripting() or is_compiling():
@@ -384,9 +328,9 @@ class MessagePassing(torch.nn.Module):
                 return src.index_select(self.node_dim, row)
 
         raise ValueError(
-            ('`MessagePassing.propagate` only supports integer tensors of '
-             'shape `[2, num_messages]`, `torch_sparse.SparseTensor` '
-             'or `torch.sparse.Tensor` for argument `edge_index`.'))
+            '`MessagePassing.propagate` only supports integer tensors of '
+            'shape `[2, num_messages]`, `torch_sparse.SparseTensor` '
+            'or `torch.sparse.Tensor` for argument `edge_index`.')
 
     def _collect(
         self,
@@ -473,7 +417,6 @@ class MessagePassing(torch.nn.Module):
 
     def forward(self, *args: Any, **kwargs: Any) -> Any:
         r"""Runs the forward pass of the module."""
-        pass
 
     def propagate(
         self,
@@ -980,6 +923,81 @@ class MessagePassing(torch.nn.Module):
 
     # TorchScript Support #####################################################
 
+    def _set_jittable_templates(self, raise_on_error: bool = False) -> None:
+        root_dir = osp.dirname(osp.realpath(__file__))
+        jinja_prefix = f'{self.__module__}_{self.__class__.__name__}'
+        # Optimize `propagate()` via `*.jinja` templates:
+        if not self.propagate.__module__.startswith(jinja_prefix):
+            try:
+                if ('propagate' in self.__class__.__dict__
+                        and self.__class__.__dict__['propagate']
+                        != MessagePassing.propagate):
+                    raise ValueError("Cannot compile custom 'propagate' "
+                                     "method")
+
+                module = module_from_template(
+                    module_name=f'{jinja_prefix}_propagate',
+                    template_path=osp.join(root_dir, 'propagate.jinja'),
+                    tmp_dirname='message_passing',
+                    # Keyword arguments:
+                    modules=self.inspector._modules,
+                    collect_name='collect',
+                    signature=self._get_propagate_signature(),
+                    collect_param_dict=self.inspector.get_flat_param_dict(
+                        ['message', 'aggregate', 'update']),
+                    message_args=self.inspector.get_param_names('message'),
+                    aggregate_args=self.inspector.get_param_names('aggregate'),
+                    message_and_aggregate_args=self.inspector.get_param_names(
+                        'message_and_aggregate'),
+                    update_args=self.inspector.get_param_names('update'),
+                    fuse=self.fuse,
+                )
+
+                self.__class__._orig_propagate = self.__class__.propagate
+                self.__class__._jinja_propagate = module.propagate
+
+                self.__class__.propagate = module.propagate
+                self.__class__.collect = module.collect
+            except Exception as e:  # pragma: no cover
+                if raise_on_error:
+                    raise e
+                self.__class__._orig_propagate = self.__class__.propagate
+                self.__class__._jinja_propagate = self.__class__.propagate
+
+        # Optimize `edge_updater()` via `*.jinja` templates (if implemented):
+        if (self.inspector.implements('edge_update')
+                and not self.edge_updater.__module__.startswith(jinja_prefix)):
+            try:
+                if ('edge_updater' in self.__class__.__dict__
+                        and self.__class__.__dict__['edge_updater']
+                        != MessagePassing.edge_updater):
+                    raise ValueError("Cannot compile custom 'edge_updater' "
+                                     "method")
+
+                module = module_from_template(
+                    module_name=f'{jinja_prefix}_edge_updater',
+                    template_path=osp.join(root_dir, 'edge_updater.jinja'),
+                    tmp_dirname='message_passing',
+                    # Keyword arguments:
+                    modules=self.inspector._modules,
+                    collect_name='edge_collect',
+                    signature=self._get_edge_updater_signature(),
+                    collect_param_dict=self.inspector.get_param_dict(
+                        'edge_update'),
+                )
+
+                self.__class__._orig_edge_updater = self.__class__.edge_updater
+                self.__class__._jinja_edge_updater = module.edge_updater
+
+                self.__class__.edge_updater = module.edge_updater
+                self.__class__.edge_collect = module.edge_collect
+            except Exception as e:  # pragma: no cover
+                if raise_on_error:
+                    raise e
+                self.__class__._orig_edge_updater = self.__class__.edge_updater
+                self.__class__._jinja_edge_updater = (
+                    self.__class__.edge_updater)
+
     def _get_propagate_signature(self) -> Signature:
         param_dict = self.inspector.get_params_from_method_call(
             'propagate', exclude=[0, 'edge_index', 'size'])

torch_geometric/nn/conv/mixhop_conv.py

@@ -14,7 +14,7 @@ from torch_geometric.utils import spmm
 
 class MixHopConv(MessagePassing):
     r"""The Mix-Hop graph convolutional operator from the
-    `"MixHop: Higher-Order Graph Convolutional Architecturesvia Sparsified
+    `"MixHop: Higher-Order Graph Convolutional Architectures via Sparsified
     Neighborhood Mixing" <https://arxiv.org/abs/1905.00067>`_ paper.
 
     .. math::

torch_geometric/nn/conv/rgcn_conv.py

@@ -120,7 +120,8 @@ class RGCNConv(MessagePassing):
             in_channels = (in_channels, in_channels)
         self.in_channels_l = in_channels[0]
 
-        self._use_segment_matmul_heuristic_output: Optional[bool] = None
+        self._use_segment_matmul_heuristic_output: torch.jit.Attribute(
+            None, Optional[float])
 
         if num_bases is not None:
             self.weight = Parameter(

torch_geometric/nn/conv/spline_conv.py

@@ -5,17 +5,17 @@ import torch
 from torch import Tensor, nn
 from torch.nn import Parameter
 
+import torch_geometric.typing
 from torch_geometric.nn.conv import MessagePassing
 from torch_geometric.nn.dense.linear import Linear
 from torch_geometric.nn.inits import uniform, zeros
 from torch_geometric.typing import Adj, OptPairTensor, OptTensor, Size
 from torch_geometric.utils.repeat import repeat
 
-try:
+if torch_geometric.typing.WITH_TORCH_SPLINE_CONV:
     from torch_spline_conv import spline_basis, spline_weighting
-except (ImportError, OSError):  # Fail gracefully on GLIBC errors
-    spline_basis = None
-    spline_weighting = None
+else:
+    spline_basis = spline_weighting = None
 
 
 class SplineConv(MessagePassing):

torch_geometric/nn/conv/x_conv.py

@@ -9,12 +9,13 @@ from torch.nn import Conv1d
 from torch.nn import Linear as L
 from torch.nn import Sequential as S
 
+import torch_geometric.typing
 from torch_geometric.nn import Reshape
 from torch_geometric.nn.inits import reset
 
-try:
+if torch_geometric.typing.WITH_TORCH_CLUSTER:
     from torch_cluster import knn_graph
-except ImportError:
+else:
     knn_graph = None
 
 

torch_geometric/nn/dense/linear.py

@@ -58,7 +58,7 @@ def reset_bias_(bias: Optional[Tensor], in_channels: int,
 
 
 class Linear(torch.nn.Module):
-    r"""Applies a linear tranformation to the incoming data.
+    r"""Applies a linear transformation to the incoming data.
 
     .. math::
         \mathbf{x}^{\prime} = \mathbf{x} \mathbf{W}^{\top} + \mathbf{b}
@@ -192,7 +192,7 @@ class Linear(torch.nn.Module):
 
 
 class HeteroLinear(torch.nn.Module):
-    r"""Applies separate linear tranformations to the incoming data according
+    r"""Applies separate linear transformations to the incoming data according
     to types.
 
     For type :math:`\kappa`, it computes
@@ -365,7 +365,8 @@ class HeteroLinear(torch.nn.Module):
 
 
 class HeteroDictLinear(torch.nn.Module):
-    r"""Applies separate linear tranformations to the incoming data dictionary.
+    r"""Applies separate linear transformations to the incoming data
+    dictionary.
 
     For key :math:`\kappa`, it computes
 
@@ -479,7 +480,7 @@ class HeteroDictLinear(torch.nn.Module):
             lin = self.lins[key]
             if is_uninitialized_parameter(lin.weight):
                 self.lins[key].initialize_parameters(None, x)
-        self.reset_parameters()
+                self.lins[key].reset_parameters()
         self._hook.remove()
         self.in_channels = {key: x.size(-1) for key, x in input[0].items()}
         delattr(self, '_hook')

torch_geometric/nn/fx.py

@@ -18,8 +18,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
 
@@ -283,7 +283,7 @@ 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

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:
@@ -175,7 +177,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()
 

torch_geometric/nn/models/__init__.py

@@ -2,7 +2,7 @@ r"""Model package."""
 
 from .mlp import MLP
 from .basic_gnn import GCN, GraphSAGE, GIN, GAT, PNA, EdgeCNN
-from .jumping_knowledge import JumpingKnowledge
+from .jumping_knowledge import JumpingKnowledge, HeteroJumpingKnowledge
 from .meta import MetaLayer
 from .node2vec import Node2Vec
 from .deep_graph_infomax import DeepGraphInfomax
@@ -28,7 +28,10 @@ from .gnnff import GNNFF
 from .pmlp import PMLP
 from .neural_fingerprint import NeuralFingerprint
 from .visnet import ViSNet
-
+from .g_retriever import GRetriever
+from .git_mol import GITMol
+from .molecule_gpt import MoleculeGPT
+from .glem import GLEM
 # Deprecated:
 from torch_geometric.explain.algorithm.captum import (to_captum_input,
                                                       captum_output_to_dicts)
@@ -42,6 +45,7 @@ __all__ = classes = [
     'PNA',
     'EdgeCNN',
     'JumpingKnowledge',
+    'HeteroJumpingKnowledge',
     'MetaLayer',
     'Node2Vec',
     'DeepGraphInfomax',
@@ -74,4 +78,8 @@ __all__ = classes = [
     'PMLP',
     'NeuralFingerprint',
     'ViSNet',
+    'GRetriever',
+    'GITMol',
+    'MoleculeGPT',
+    'GLEM',
 ]

torch_geometric/nn/models/deep_graph_infomax.py

@@ -98,7 +98,6 @@ class DeepGraphInfomax(torch.nn.Module):
         test_z: Tensor,
         test_y: Tensor,
         solver: str = 'lbfgs',
-        multi_class: str = 'auto',
         *args,
         **kwargs,
     ) -> float:
@@ -107,7 +106,7 @@ class DeepGraphInfomax(torch.nn.Module):
         """
         from sklearn.linear_model import LogisticRegression
 
-        clf = LogisticRegression(solver=solver, multi_class=multi_class, *args,
+        clf = LogisticRegression(solver=solver, *args,
                                  **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_utils.py

@@ -1,14 +1,12 @@
 # Shameless steal from: https://github.com/klicperajo/dimenet
 
+import math
+
 import numpy as np
+import sympy as sym
 from scipy import special as sp
 from scipy.optimize import brentq
 
-try:
-    import sympy as sym
-except ImportError:
-    sym = None
-
 
 def Jn(r, n):
     return np.sqrt(np.pi / (2 * r)) * sp.jv(n + 0.5, r)
@@ -66,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):