pyg-nightly 2.6.0.dev20240318__py3-none-any.whl → 2.7.0.dev20250115__py3-none-any.whl

torch_geometric/nn/conv/hetero_conv.py

@@ -70,8 +70,8 @@ class HeteroConv(torch.nn.Module):
         for edge_type, module in convs.items():
             check_add_self_loops(module, [edge_type])
 
-        src_node_types = set([key[0] for key in convs.keys()])
-        dst_node_types = set([key[-1] for key in convs.keys()])
+        src_node_types = {key[0] for key in convs.keys()}
+        dst_node_types = {key[-1] for key in convs.keys()}
         if len(src_node_types - dst_node_types) > 0:
             warnings.warn(
                 f"There exist node types ({src_node_types - dst_node_types}) "
@@ -102,7 +102,7 @@ class HeteroConv(torch.nn.Module):
                 individual edge type, either as a :class:`torch.Tensor` of
                 shape :obj:`[2, num_edges]` or a
                 :class:`torch_sparse.SparseTensor`.
-            *args_dict (optional): Additional forward arguments of invididual
+            *args_dict (optional): Additional forward arguments of individual
                 :class:`torch_geometric.nn.conv.MessagePassing` layers.
             **kwargs_dict (optional): Additional forward arguments of
                 individual :class:`torch_geometric.nn.conv.MessagePassing`

torch_geometric/nn/conv/hgt_conv.py

@@ -67,7 +67,7 @@ class HGTConv(MessagePassing):
             for i, edge_type in enumerate(metadata[1])
         }
 
-        self.dst_node_types = set([key[-1] for key in self.edge_types])
+        self.dst_node_types = {key[-1] for key in self.edge_types}
 
         self.kqv_lin = HeteroDictLinear(self.in_channels,
                                         self.out_channels * 3)

torch_geometric/nn/conv/message_passing.py

@@ -6,6 +6,7 @@ from typing import (
     Any,
     Callable,
     Dict,
+    Final,
     List,
     Optional,
     OrderedDict,
@@ -19,6 +20,7 @@ from torch import Tensor
 from torch.utils.hooks import RemovableHandle
 
 from torch_geometric import EdgeIndex, is_compiling
+from torch_geometric.index import ptr2index
 from torch_geometric.inspector import Inspector, Signature
 from torch_geometric.nn.aggr import Aggregation
 from torch_geometric.nn.resolver import aggregation_resolver as aggr_resolver
@@ -29,7 +31,6 @@ from torch_geometric.utils import (
     is_torch_sparse_tensor,
     to_edge_index,
 )
-from torch_geometric.utils.sparse import ptr2index
 
 FUSE_AGGRS = {'add', 'sum', 'mean', 'min', 'max'}
 HookDict = OrderedDict[int, Callable]
@@ -102,6 +103,10 @@ class MessagePassing(torch.nn.Module):
         'size_i', 'size_j', 'ptr', 'index', 'dim_size'
     }
 
+    # Supports `message_and_aggregate` via `EdgeIndex`.
+    # TODO Remove once migration is finished.
+    SUPPORTS_FUSED_EDGE_INDEX: Final[bool] = False
+
     def __init__(
         self,
         aggr: Optional[Union[str, List[str], Aggregation]] = 'sum',
@@ -162,63 +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:
-                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:
-                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
@@ -227,6 +177,7 @@ class MessagePassing(torch.nn.Module):
         self._apply_sigmoid: bool = True
 
         # Inference Decomposition:
+        self._decomposed_layers = 1
         self.decomposed_layers = decomposed_layers
 
     def reset_parameters(self) -> None:
@@ -234,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'):
@@ -247,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):
@@ -256,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 "
@@ -284,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,
@@ -299,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():
@@ -370,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,
@@ -459,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,
@@ -509,7 +466,17 @@ class MessagePassing(torch.nn.Module):
         mutable_size = self._check_input(edge_index, size)
 
         # Run "fused" message and aggregation (if applicable).
-        if is_sparse(edge_index) and self.fuse and not self.explain:
+        fuse = False
+        if self.fuse and not self.explain:
+            if is_sparse(edge_index):
+                fuse = True
+            elif (not torch.jit.is_scripting()
+                  and isinstance(edge_index, EdgeIndex)):
+                if (self.SUPPORTS_FUSED_EDGE_INDEX
+                        and edge_index.is_sorted_by_col):
+                    fuse = True
+
+        if fuse:
             coll_dict = self._collect(self._fused_user_args, edge_index,
                                       mutable_size, kwargs)
 
@@ -628,7 +595,7 @@ class MessagePassing(torch.nn.Module):
                                 dim=self.node_dim)
 
     @abstractmethod
-    def message_and_aggregate(self, adj_t: Adj) -> Tensor:
+    def message_and_aggregate(self, edge_index: Adj) -> Tensor:
         r"""Fuses computations of :func:`message` and :func:`aggregate` into a
         single function.
         If applicable, this saves both time and memory since messages do not
@@ -720,16 +687,20 @@ class MessagePassing(torch.nn.Module):
             raise ValueError("Inference decomposition of message passing "
                              "modules is only supported on the Python module")
 
+        if decomposed_layers == self._decomposed_layers:
+            return  # Abort early if nothing to do.
+
         self._decomposed_layers = decomposed_layers
 
         if decomposed_layers != 1:
-            self.propagate = self.__class__._orig_propagate.__get__(
-                self, MessagePassing)
+            if hasattr(self.__class__, '_orig_propagate'):
+                self.propagate = self.__class__._orig_propagate.__get__(
+                    self, MessagePassing)
 
-        elif ((self.explain is None or self.explain is False)
-              and not self.propagate.__module__.endswith('_propagate')):
-            self.propagate = self.__class__._jinja_propagate.__get__(
-                self, MessagePassing)
+        elif self.explain is None or self.explain is False:
+            if hasattr(self.__class__, '_jinja_propagate'):
+                self.propagate = self.__class__._jinja_propagate.__get__(
+                    self, MessagePassing)
 
     # Explainability ##########################################################
 
@@ -743,6 +714,9 @@ class MessagePassing(torch.nn.Module):
             raise ValueError("Explainability of message passing modules "
                              "is only supported on the Python module")
 
+        if explain == self._explain:
+            return  # Abort early if nothing to do.
+
         self._explain = explain
 
         if explain is True:
@@ -753,16 +727,18 @@ class MessagePassing(torch.nn.Module):
                 funcs=['message', 'explain_message', 'aggregate', 'update'],
                 exclude=self.special_args,
             )
-            self.propagate = self.__class__._orig_propagate.__get__(
-                self, MessagePassing)
+            if hasattr(self.__class__, '_orig_propagate'):
+                self.propagate = self.__class__._orig_propagate.__get__(
+                    self, MessagePassing)
         else:
             self._user_args = self.inspector.get_flat_param_names(
                 funcs=['message', 'aggregate', 'update'],
                 exclude=self.special_args,
             )
             if self.decomposed_layers == 1:
-                self.propagate = self.__class__._jinja_propagate.__get__(
-                    self, MessagePassing)
+                if hasattr(self.__class__, '_jinja_propagate'):
+                    self.propagate = self.__class__._jinja_propagate.__get__(
+                        self, MessagePassing)
 
     def explain_message(
         self,
@@ -947,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/propagate.jinja

@@ -42,7 +42,15 @@ def propagate(
     # End Propagate Forward Pre Hook ###########################################
 
     mutable_size = self._check_input(edge_index, size)
-    fuse = is_sparse(edge_index) and self.fuse
+
+    # Run "fused" message and aggregation (if applicable).
+    fuse = False
+    if self.fuse:
+        if is_sparse(edge_index):
+            fuse = True
+        elif not torch.jit.is_scripting() and isinstance(edge_index, EdgeIndex):
+            if self.SUPPORTS_FUSED_EDGE_INDEX and edge_index.is_sorted_by_col:
+                fuse = True
 
     if fuse:
 

torch_geometric/nn/conv/rgcn_conv.py

@@ -3,11 +3,11 @@ from typing import Optional, Tuple, Union
 import torch
 from torch import Tensor
 from torch.nn import Parameter
-from torch.nn import Parameter as Param
 
 import torch_geometric.backend
 import torch_geometric.typing
 from torch_geometric import is_compiling
+from torch_geometric.index import index2ptr
 from torch_geometric.nn.conv import MessagePassing
 from torch_geometric.nn.inits import glorot, zeros
 from torch_geometric.typing import (
@@ -18,7 +18,6 @@ from torch_geometric.typing import (
     torch_sparse,
 )
 from torch_geometric.utils import index_sort, one_hot, scatter, spmm
-from torch_geometric.utils.sparse import index2ptr
 
 
 def masked_edge_index(edge_index: Adj, edge_mask: Tensor) -> Adj:
@@ -121,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(
@@ -143,12 +143,12 @@ class RGCNConv(MessagePassing):
             self.register_parameter('comp', None)
 
         if root_weight:
-            self.root = Param(torch.empty(in_channels[1], out_channels))
+            self.root = Parameter(torch.empty(in_channels[1], out_channels))
         else:
             self.register_parameter('root', None)
 
         if bias:
-            self.bias = Param(torch.empty(out_channels))
+            self.bias = Parameter(torch.empty(out_channels))
         else:
             self.register_parameter('bias', None)
 

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

@@ -12,10 +12,10 @@ from torch.nn.parameter import Parameter
 import torch_geometric.backend
 import torch_geometric.typing
 from torch_geometric import is_compiling
+from torch_geometric.index import index2ptr
 from torch_geometric.nn import inits
 from torch_geometric.typing import pyg_lib
 from torch_geometric.utils import index_sort
-from torch_geometric.utils.sparse import index2ptr
 
 
 def is_uninitialized_parameter(x: Any) -> bool:
@@ -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
@@ -222,6 +222,8 @@ class HeteroLinear(torch.nn.Module):
           type vector :math:`(*)`
         - **output:** features :math:`(*, F_{out})`
     """
+    _timing_cache: Dict[int, Tuple[float, float]]
+
     def __init__(
         self,
         in_channels: int,
@@ -245,15 +247,17 @@ class HeteroLinear(torch.nn.Module):
         else:
             self.weight = torch.nn.Parameter(
                 torch.empty(num_types, in_channels, out_channels))
+
         if kwargs.get('bias', True):
             self.bias = Parameter(torch.empty(num_types, out_channels))
         else:
             self.register_parameter('bias', None)
-        self.reset_parameters()
 
         # Timing cache for benchmarking naive vs. segment matmul usage:
         self._timing_cache: Dict[int, Tuple[float, float]] = {}
 
+        self.reset_parameters()
+
     def reset_parameters(self):
         r"""Resets all learnable parameters of the module."""
         reset_weight_(self.weight, self.in_channels,
@@ -361,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
 
@@ -475,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',
 ]