Repository: pyg-team/pytorch_geometric
Pull request #9499 (2024-07-09 14:40:05)

Fix `visualize_feature_importance` in heterogenous graphs with single features
When calling visualize_feature_importance for a HeteroExplanation object, an error occurs if one of the Heterogenous graph's node types has one single feature. The error message also had a bug, as "node_mask_dict" has no attribute "size()". 

<img width="591" alt="image" src="https://github.com/pyg-team/pytorch_geometric/assets/160508240/f078882a-79cf-40e0-8162-4afb20df1b13">


This is a simple fix for both bugs.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4d9a45259..0f564a8f4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added the `PatchTransformerAggregation` layer ([#9487](https://github.com/pyg-team/pytorch_geometric/pull/9487))
 - Added the `nn.nlp.LLM` model ([#9462](https://github.com/pyg-team/pytorch_geometric/pull/9462))
 - Added an example of training GNNs for a graph-level regression task ([#9070](https://github.com/pyg-team/pytorch_geometric/pull/9070))
 - Added `utils.from_rdmol`/`utils.to_rdmol` functionality ([#9452](https://github.com/pyg-team/pytorch_geometric/pull/9452))
diff --git a/examples/triangles_sag_pool.py b/examples/triangles_sag_pool.py
index a3d3898f7..8a6831d5c 100644
--- a/examples/triangles_sag_pool.py
+++ b/examples/triangles_sag_pool.py
@@ -71,7 +71,7 @@ model = Net(dataset.num_features).to(device)
 optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
 
 
-def train(epoch):
+def train():
     model.train()
 
     total_loss = 0
@@ -101,7 +101,7 @@ def test(loader):
 
 
 for epoch in range(1, 301):
-    loss = train(epoch)
+    loss = train()
     train_correct, train_ratio = test(train_loader)
     val_correct, val_ratio = test(val_loader)
     test_correct, test_ratio = test(test_loader)
diff --git a/torch_geometric/config_mixin.py b/torch_geometric/config_mixin.py
index 0cb89fab8..ab0b5b5ff 100644
--- a/torch_geometric/config_mixin.py
+++ b/torch_geometric/config_mixin.py
@@ -82,7 +82,7 @@ def _recursive_from_config(value: Any) -> Any:
     cls: Any = None
     if is_dataclass(value):
         if getattr(value, '_target_', None):
-            cls = _locate_cls(value._target_)
+            cls = _locate_cls(value._target_)  # type: ignore[attr-defined]
         else:
             cls = class_from_dataclass(value.__class__)
     elif isinstance(value, dict) and '_target_' in value:
diff --git a/torch_geometric/edge_index.py b/torch_geometric/edge_index.py
index 17ee648fe..3029d78ea 100644
--- a/torch_geometric/edge_index.py
+++ b/torch_geometric/edge_index.py
@@ -1928,7 +1928,7 @@ def _spmm(
     if transpose and not input.is_sorted_by_col:
         cls_name = input.__class__.__name__
         raise ValueError(f"'matmul(..., transpose=True)' requires "
-                         f"'{cls_name}' to be sorted by colums")
+                         f"'{cls_name}' to be sorted by columns")
 
     if (torch_geometric.typing.WITH_TORCH_SPARSE and not is_compiling()
             and other.is_cuda):  # pragma: no cover
diff --git a/torch_geometric/explain/explanation.py b/torch_geometric/explain/explanation.py
index 72ef79bde..2aca9b280 100644
--- a/torch_geometric/explain/explanation.py
+++ b/torch_geometric/explain/explanation.py
@@ -340,10 +340,10 @@ class HeteroExplanation(HeteroData, ExplanationMixin):
         """
         node_mask_dict = self.node_mask_dict
         for node_mask in node_mask_dict.values():
-            if node_mask.dim() != 2 or node_mask.size(1) <= 1:
+            if node_mask.dim() != 2:
                 raise ValueError(f"Cannot compute feature importance for "
                                  f"object-level 'node_mask' "
-                                 f"(got shape {node_mask_dict.size()})")
+                                 f"(got shape {node_mask.size()})")
 
         if feat_labels is None:
             feat_labels = {}
diff --git a/torch_geometric/inspector.py b/torch_geometric/inspector.py
index 8f6e7bd04..32b4421f7 100644
--- a/torch_geometric/inspector.py
+++ b/torch_geometric/inspector.py
@@ -305,7 +305,7 @@ class Inspector:
         according to its function signature from a data blob.
 
         Args:
-            func (callabel or str): The function.
+            func (callable or str): The function.
             kwargs (dict[str, Any]): The data blob which may serve as inputs.
         """
         out_dict: Dict[str, Any] = {}
@@ -346,7 +346,7 @@ class Inspector:
             type annotations are not found.
 
         Args:
-            func (callabel or str): The function.
+            func (callable or str): The function.
             exclude (list[int or str]): A list of parameters to exclude, either
                 given by their name or index. (default: :obj:`None`)
         """
diff --git a/torch_geometric/nn/aggr/__init__.py b/torch_geometric/nn/aggr/__init__.py
index 0b3425849..aaf8c95e9 100644
--- a/torch_geometric/nn/aggr/__init__.py
+++ b/torch_geometric/nn/aggr/__init__.py
@@ -25,6 +25,7 @@ from .deep_sets import DeepSetsAggregation
 from .set_transformer import SetTransformerAggregation
 from .lcm import LCMAggregation
 from .variance_preserving import VariancePreservingAggregation
+from .patch_transformer import PatchTransformerAggregation
 
 __all__ = classes = [
     'Aggregation',
@@ -53,4 +54,5 @@ __all__ = classes = [
     'SetTransformerAggregation',
     'LCMAggregation',
     'VariancePreservingAggregation',
+    'PatchTransformerAggregation',
 ]
diff --git a/torch_geometric/nn/aggr/patch_transformer.py b/torch_geometric/nn/aggr/patch_transformer.py
new file mode 100644
index 000000000..480408488
--- /dev/null
+++ b/torch_geometric/nn/aggr/patch_transformer.py
@@ -0,0 +1,143 @@
+import math
+from typing import List, Optional, Union
+
+import torch
+from torch import Tensor
+
+from torch_geometric.experimental import disable_dynamic_shapes
+from torch_geometric.nn.aggr import Aggregation
+from torch_geometric.nn.aggr.utils import MultiheadAttentionBlock
+from torch_geometric.nn.encoding import PositionalEncoding
+from torch_geometric.utils import scatter
+
+
+class PatchTransformerAggregation(Aggregation):
+    r"""Performs patch transformer aggregation in which the elements to
+    aggregate are processed by multi-head attention blocks across patches, as
+    described in the `"Simplifying Temporal Heterogeneous Network for
+    Continuous-Time Link Prediction"
+    <https://dl.acm.org/doi/pdf/10.1145/3583780.3615059>`_ paper.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        out_channels (int): Size of each output sample.
+        patch_size (int): Number of elements in a patch.
+        hidden_channels (int): Intermediate size of each sample.
+        num_transformer_blocks (int, optional): Number of transformer blocks
+            (default: :obj:`1`).
+        heads (int, optional): Number of multi-head-attentions.
+            (default: :obj:`1`)
+        dropout (float, optional): Dropout probability of attention weights.
+            (default: :obj:`0.0`)
+        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"`)
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        patch_size: int,
+        hidden_channels: int,
+        num_transformer_blocks: int = 1,
+        heads: int = 1,
+        dropout: float = 0.0,
+        aggr: Union[str, List[str]] = 'mean',
+    ) -> None:
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.patch_size = patch_size
+        self.aggrs = [aggr] if isinstance(aggr, str) else aggr
+
+        assert len(self.aggrs) > 0
+        for aggr in self.aggrs:
+            assert aggr in ['sum', 'mean', 'min', 'max', 'var', 'std']
+
+        self.lin = torch.nn.Linear(in_channels, hidden_channels)
+        self.pad_projector = torch.nn.Linear(
+            patch_size * hidden_channels,
+            hidden_channels,
+        )
+        self.pe = PositionalEncoding(hidden_channels)
+
+        self.blocks = torch.nn.ModuleList([
+            MultiheadAttentionBlock(
+                channels=hidden_channels,
+                heads=heads,
+                layer_norm=True,
+                dropout=dropout,
+            ) for _ in range(num_transformer_blocks)
+        ])
+
+        self.fc = torch.nn.Linear(
+            hidden_channels * len(self.aggrs),
+            out_channels,
+        )
+
+    def reset_parameters(self) -> None:
+        self.lin.reset_parameters()
+        self.pad_projector.reset_parameters()
+        self.pe.reset_parameters()
+        for block in self.blocks:
+            block.reset_parameters()
+        self.fc.reset_parameters()
+
+    @disable_dynamic_shapes(required_args=['dim_size', 'max_num_elements'])
+    def forward(
+        self,
+        x: Tensor,
+        index: Tensor,
+        ptr: Optional[Tensor] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        max_num_elements: Optional[int] = None,
+    ) -> Tensor:
+
+        if max_num_elements is None:
+            if ptr is not None:
+                count = ptr.diff()
+            else:
+                count = scatter(torch.ones_like(index), index, dim=0,
+                                dim_size=dim_size, reduce='sum')
+            max_num_elements = int(count.max()) + 1
+
+        # Set `max_num_elements` to a multiple of `patch_size`:
+        max_num_elements = (math.floor(max_num_elements / self.patch_size) *
+                            self.patch_size)
+
+        x = self.lin(x)
+
+        # TODO If groups are heavily unbalanced, this will create a lot of
+        # "empty" patches. Try to figure out a way to fix this.
+        # [batch_size, num_patches * patch_size, hidden_channels]
+        x, _ = self.to_dense_batch(x, index, ptr, dim_size, dim,
+                                   max_num_elements=max_num_elements)
+
+        # [batch_size, num_patches, patch_size * hidden_channels]
+        x = x.view(x.size(0), max_num_elements // self.patch_size,
+                   self.patch_size * x.size(-1))
+
+        # [batch_size, num_patches, hidden_channels]
+        x = self.pad_projector(x)
+
+        x = x + self.pe(torch.arange(x.size(1), device=x.device))
+
+        # [batch_size, num_patches, hidden_channels]
+        for block in self.blocks:
+            x = block(x, x)
+
+        # [batch_size, hidden_channels]
+        outs: List[Tensor] = []
+        for aggr in self.aggrs:
+            out = getattr(torch, aggr)(x, dim=1)
+            outs.append(out[0] if isinstance(out, tuple) else out)
+        out = torch.cat(outs, dim=1) if len(outs) > 1 else outs[0]
+
+        # [batch_size, out_channels]
+        return self.fc(out)
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({self.in_channels}, '
+                f'{self.out_channels}, patch_size={self.patch_size})')
diff --git a/torch_geometric/nn/conv/hetero_conv.py b/torch_geometric/nn/conv/hetero_conv.py
index 2975854f0..29c6150ae 100644
--- a/torch_geometric/nn/conv/hetero_conv.py
+++ b/torch_geometric/nn/conv/hetero_conv.py
@@ -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`
diff --git a/torch_geometric/nn/conv/mixhop_conv.py b/torch_geometric/nn/conv/mixhop_conv.py
index c8e0be466..bc51a8550 100644
--- a/torch_geometric/nn/conv/mixhop_conv.py
+++ b/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::
diff --git a/torch_geometric/nn/dense/linear.py b/torch_geometric/nn/dense/linear.py
index 6225fde7d..564b13161 100644
--- a/torch_geometric/nn/dense/linear.py
+++ b/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
 
diff --git a/torch_geometric/nn/fx.py b/torch_geometric/nn/fx.py
index d1141e44c..6a266bdd5 100644
--- a/torch_geometric/nn/fx.py
+++ b/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
diff --git a/torch_geometric/nn/nlp/sentence_transformer.py b/torch_geometric/nn/nlp/sentence_transformer.py
index cfb26176f..b1ae8254c 100644
--- a/torch_geometric/nn/nlp/sentence_transformer.py
+++ b/torch_geometric/nn/nlp/sentence_transformer.py
@@ -27,6 +27,8 @@ class SentenceTransformer(torch.nn.Module):
 
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
         self.model = AutoModel.from_pretrained(model_name)
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
 
     def forward(self, input_ids: Tensor, attention_mask: Tensor) -> Tensor:
         out = self.model(input_ids=input_ids, attention_mask=attention_mask)
diff --git a/torch_geometric/nn/summary.py b/torch_geometric/nn/summary.py
index 8256bacc9..a8aa2e472 100644
--- a/torch_geometric/nn/summary.py
+++ b/torch_geometric/nn/summary.py
@@ -141,7 +141,7 @@ def get_shape(inputs: Any) -> str:
 def postprocess(info_list: List[dict]) -> List[dict]:
     for idx, info in enumerate(info_list):
         depth = info['depth']
-        if idx > 0:  # root module (0) is exclued
+        if idx > 0:  # root module (0) is excluded
             if depth == 1:
                 prefix = '├─'
             else:
diff --git a/torch_geometric/nn/to_hetero_with_bases_transformer.py b/torch_geometric/nn/to_hetero_with_bases_transformer.py
index b99ee9e9f..5d2526b34 100644
--- a/torch_geometric/nn/to_hetero_with_bases_transformer.py
+++ b/torch_geometric/nn/to_hetero_with_bases_transformer.py
@@ -468,7 +468,7 @@ def get_edge_type(
 ###############################################################################
 
 # These methods are used to group the individual type-wise components into a
-# unfied single representation.
+# unified single representation.
 
 
 def group_node_placeholder(input_dict: Dict[NodeType, Tensor],
diff --git a/torch_geometric/sampler/base.py b/torch_geometric/sampler/base.py
index 8383af2e7..d67ddd5af 100644
--- a/torch_geometric/sampler/base.py
+++ b/torch_geometric/sampler/base.py
@@ -537,7 +537,7 @@ class NegativeSampling(CastMixin):
             the sampling of source nodes. Does not necessarily need to sum up
             to one. If not given, negative nodes will be sampled uniformly.
             (default: :obj:`None`)
-        src_weight (torch.Tensor, optional): A node-level vector determining
+        dst_weight (torch.Tensor, optional): A node-level vector determining
             the sampling of destination nodes. Does not necessarily need to sum
             up to one. If not given, negative nodes will be sampled uniformly.
             (default: :obj:`None`)
diff --git a/torch_geometric/sampler/neighbor_sampler.py b/torch_geometric/sampler/neighbor_sampler.py
index 758f14819..0ce692206 100644
--- a/torch_geometric/sampler/neighbor_sampler.py
+++ b/torch_geometric/sampler/neighbor_sampler.py
@@ -631,7 +631,7 @@ def edge_sample(
             if edge_label_time is not None:
                 dst_time = edge_label_time.repeat(1 + neg_sampling.amount)
 
-    # Heterogeneus Neighborhood Sampling ######################################
+    # Heterogeneous Neighborhood Sampling #####################################
 
     if input_type is not None:
         seed_time_dict = None
@@ -724,7 +724,7 @@ def edge_sample(
                 src_time,
             )
 
-    # Homogeneus Neighborhood Sampling ########################################
+    # Homogeneous Neighborhood Sampling #######################################
 
     else:
 
diff --git a/torch_geometric/transforms/add_metapaths.py b/torch_geometric/transforms/add_metapaths.py
index 9aca179fc..366a09644 100644
--- a/torch_geometric/transforms/add_metapaths.py
+++ b/torch_geometric/transforms/add_metapaths.py
@@ -37,7 +37,7 @@ class AddMetaPaths(BaseTransform):
     :class:`~torch_geometric.data.HeteroData` object as edge type
     :obj:`(src_node_type, "metapath_*", dst_node_type)`, where
     :obj:`src_node_type` and :obj:`dst_node_type` denote :math:`\mathcal{V}_1`
-    and :math:`\mathcal{V}_{\ell}`, repectively.
+    and :math:`\mathcal{V}_{\ell}`, respectively.
 
     In addition, a :obj:`metapath_dict` object is added to the
     :class:`~torch_geometric.data.HeteroData` object which maps the
@@ -108,12 +108,12 @@ class AddMetaPaths(BaseTransform):
         **kwargs: bool,
     ) -> None:
         if 'drop_orig_edges' in kwargs:
-            warnings.warn("'drop_orig_edges' is dprecated. Use "
+            warnings.warn("'drop_orig_edges' is deprecated. Use "
                           "'drop_orig_edge_types' instead")
             drop_orig_edge_types = kwargs['drop_orig_edges']
 
         if 'drop_unconnected_nodes' in kwargs:
-            warnings.warn("'drop_unconnected_nodes' is dprecated. Use "
+            warnings.warn("'drop_unconnected_nodes' is deprecated. Use "
                           "'drop_unconnected_node_types' instead")
             drop_unconnected_node_types = kwargs['drop_unconnected_nodes']
 
@@ -231,7 +231,7 @@ class AddRandomMetaPaths(BaseTransform):
             will drop node types not connected by any edge type.
             (default: :obj:`False`)
         walks_per_node (int, List[int], optional): The number of random walks
-            for each starting node in a metapth. (default: :obj:`1`)
+            for each starting node in a metapath. (default: :obj:`1`)
         sample_ratio (float, optional): The ratio of source nodes to start
             random walks from. (default: :obj:`1.0`)
     """
diff --git a/torch_geometric/transforms/node_property_split.py b/torch_geometric/transforms/node_property_split.py
index 1850c7a00..ead02249b 100644
--- a/torch_geometric/transforms/node_property_split.py
+++ b/torch_geometric/transforms/node_property_split.py
@@ -53,7 +53,7 @@ class NodePropertySplit(BaseTransform):
 
         property_name = 'popularity'
         ratios = [0.3, 0.1, 0.1, 0.3, 0.2]
-        tranaform = NodePropertySplit(property_name, ratios)
+        transform = NodePropertySplit(property_name, ratios)
 
         data = transform(data)
     """
diff --git a/torch_geometric/transforms/pad.py b/torch_geometric/transforms/pad.py
index 74e08b50e..90f5c6222 100644
--- a/torch_geometric/transforms/pad.py
+++ b/torch_geometric/transforms/pad.py
@@ -262,7 +262,7 @@ class Pad(BaseTransform):
     All the attributes of node types other than :obj:`v0` and :obj:`v1` are
     padded using a value of :obj:`1.0`.
     All the attributes of the :obj:`('v0', 'e0', 'v1')` edge type are padded
-    usin a value of :obj:`3.5`.
+    using a value of :obj:`3.5`.
     The :obj:`edge_attr` attributes of the :obj:`('v1', 'e0', 'v0')` edge type
     are padded using a value of :obj:`-1.5`, and any other attributes of this
     edge type are padded using a value of :obj:`5.5`.
diff --git a/torch_geometric/transforms/svd_feature_reduction.py b/torch_geometric/transforms/svd_feature_reduction.py
index 624bb5e74..3f65982cf 100644
--- a/torch_geometric/transforms/svd_feature_reduction.py
+++ b/torch_geometric/transforms/svd_feature_reduction.py
@@ -11,7 +11,7 @@ class SVDFeatureReduction(BaseTransform):
     Decomposition (SVD) (functional name: :obj:`svd_feature_reduction`).
 
     Args:
-        out_channels (int): The dimensionlity of node features after
+        out_channels (int): The dimensionality of node features after
             reduction.
     """
     def __init__(self, out_channels: int):
diff --git a/torch_geometric/utils/_negative_sampling.py b/torch_geometric/utils/_negative_sampling.py
index 3deb6bcec..bb93a9385 100644
--- a/torch_geometric/utils/_negative_sampling.py
+++ b/torch_geometric/utils/_negative_sampling.py
@@ -265,7 +265,7 @@ def structured_negative_sampling_feasible(
     :meth:`~torch_geometric.utils.structured_negative_sampling` is feasible
     on the graph given by :obj:`edge_index`.
     :meth:`~torch_geometric.utils.structured_negative_sampling` is infeasible
-    if atleast one node is connected to all other nodes.
+    if at least one node is connected to all other nodes.
 
     Args:
         edge_index (LongTensor): The edge indices.
diff --git a/torch_geometric/utils/_scatter.py b/torch_geometric/utils/_scatter.py
index b4c8f5182..3d02d16a0 100644
--- a/torch_geometric/utils/_scatter.py
+++ b/torch_geometric/utils/_scatter.py
@@ -295,7 +295,7 @@ def group_cat(
     r"""Concatenates the given sequence of tensors :obj:`tensors` in the given
     dimension :obj:`dim`.
     Different from :meth:`torch.cat`, values along the concatenating dimension
-    are grouped according to the indicies defined in the :obj:`index` tensors.
+    are grouped according to the indices defined in the :obj:`index` tensors.
     All tensors must have the same shape (except in the concatenating
     dimension).
 
diff --git a/torch_geometric/utils/_tree_decomposition.py b/torch_geometric/utils/_tree_decomposition.py
index 5f6a7cf0d..14018e569 100644
--- a/torch_geometric/utils/_tree_decomposition.py
+++ b/torch_geometric/utils/_tree_decomposition.py
@@ -64,7 +64,7 @@ def tree_decomposition(
             xs.append(1)
 
     # Generate `atom2cliques` mappings.
-    atom2cliques: List[List[int]] = [[] for i in range(mol.GetNumAtoms())]
+    atom2cliques: List[List[int]] = [[] for _ in range(mol.GetNumAtoms())]
     for c in range(len(cliques)):
         for atom in cliques[c]:
             atom2cliques[atom].append(c)
diff --git a/torch_geometric/utils/augmentation.py b/torch_geometric/utils/augmentation.py
index 63f61d7bc..354fdb3b2 100644
--- a/torch_geometric/utils/augmentation.py
+++ b/torch_geometric/utils/augmentation.py
@@ -12,7 +12,7 @@ def shuffle_node(
     training: bool = True,
 ) -> Tuple[Tensor, Tensor]:
     r"""Randomly shuffle the feature matrix :obj:`x` along the
-    first dimmension.
+    first dimension.
 
     The method returns (1) the shuffled :obj:`x`, (2) the permutation
     indicating the orders of original nodes after shuffling.
diff --git a/torch_geometric/utils/map.py b/torch_geometric/utils/map.py
index 16277e6a8..554bf6109 100644
--- a/torch_geometric/utils/map.py
+++ b/torch_geometric/utils/map.py
@@ -14,7 +14,7 @@ def map_index(
     inclusive: bool = False,
 ) -> Tuple[Tensor, Optional[Tensor]]:
     r"""Maps indices in :obj:`src` to the positional value of their
-    corresponding occurence in :obj:`index`.
+    corresponding occurrence in :obj:`index`.
     Indices must be strictly positive.
 
     Args:


Test patch:
diff --git a/test/data/lightning/test_datamodule.py b/test/data/lightning/test_datamodule.py
index 81d9b9866..799cdaa3a 100644
--- a/test/data/lightning/test_datamodule.py
+++ b/test/data/lightning/test_datamodule.py
@@ -276,11 +276,19 @@ class LinearHeteroNodeModule(LightningModule):
         return torch.optim.Adam(self.parameters(), lr=0.01)
 
 
+@pytest.fixture
+def preserve_context():
+    num_threads = torch.get_num_threads()
+    yield
+    torch.distributed.destroy_process_group()
+    torch.set_num_threads(num_threads)
+
+
 @onlyCUDA
 @onlyFullTest
 @onlyNeighborSampler
 @withPackage('pytorch_lightning>=2.0.0', 'torchmetrics>=0.11.0')
-def test_lightning_hetero_node_data(get_dataset):
+def test_lightning_hetero_node_data(preserve_context, get_dataset):
     import pytorch_lightning as pl
 
     data = get_dataset(name='hetero')[0]
diff --git a/test/graphgym/test_logger.py b/test/graphgym/test_logger.py
index 74dd0c33d..d5f87ef0d 100644
--- a/test/graphgym/test_logger.py
+++ b/test/graphgym/test_logger.py
@@ -1,9 +1,15 @@
+from torch_geometric.graphgym.config import set_run_dir
+from torch_geometric.graphgym.loader import create_loader
 from torch_geometric.graphgym.logger import Logger, LoggerCallback
 from torch_geometric.testing import withPackage
 
 
 @withPackage('yacs', 'pytorch_lightning')
 def test_logger_callback():
+    loaders = create_loader()
+    assert len(loaders) == 3
+
+    set_run_dir('.')
     logger = LoggerCallback()
     assert isinstance(logger.train_logger, Logger)
     assert isinstance(logger.val_logger, Logger)
diff --git a/test/nn/aggr/test_patch_transformer.py b/test/nn/aggr/test_patch_transformer.py
new file mode 100644
index 000000000..430531048
--- /dev/null
+++ b/test/nn/aggr/test_patch_transformer.py
@@ -0,0 +1,27 @@
+import torch
+
+from torch_geometric.nn import PatchTransformerAggregation
+from torch_geometric.testing import withCUDA
+
+
+@withCUDA
+def test_patch_transformer_aggregation(device: torch.device) -> None:
+    aggr = PatchTransformerAggregation(
+        in_channels=16,
+        out_channels=32,
+        patch_size=2,
+        hidden_channels=8,
+        num_transformer_blocks=1,
+        heads=2,
+        dropout=0.2,
+        aggr=['sum', 'mean', 'min', 'max', 'var', 'std'],
+    ).to(device)
+    aggr.reset_parameters()
+    assert str(aggr) == 'PatchTransformerAggregation(16, 32, patch_size=2)'
+
+    index = torch.tensor([0, 0, 1, 1, 1, 2], device=device)
+    x = torch.randn(index.size(0), 16, device=device)
+
+    out = aggr(x, index)
+    assert out.device == device
+    assert out.size() == (3, aggr.out_channels)
diff --git a/test/utils/test_geodesic.py b/test/utils/test_geodesic.py
index c1daf41e2..de31eb512 100644
--- a/test/utils/test_geodesic.py
+++ b/test/utils/test_geodesic.py
@@ -1,6 +1,5 @@
 from math import sqrt
 
-import pytest
 import torch
 
 from torch_geometric.testing import withPackage
@@ -8,7 +7,6 @@ from torch_geometric.utils import geodesic_distance
 
 
 @withPackage('gdist')
-@pytest.mark.skip(reason="No way of currently testing this")
 def test_geodesic_distance():
     pos = torch.tensor([
         [0.0, 0.0, 0.0],
@@ -29,12 +27,12 @@ def test_geodesic_distance():
     assert torch.allclose(out, geodesic_distance(pos, face, num_workers=-1))
 
     out = geodesic_distance(pos, face, norm=False)
-    expected = [
+    expected = torch.tensor([
         [0, 2, 2, 2 * sqrt(2)],
         [2, 0, 2 * sqrt(2), 2],
         [2, 2 * sqrt(2), 0, 2],
         [2 * sqrt(2), 2, 2, 0],
-    ]
+    ])
     assert torch.allclose(out, expected)
 
     src = torch.tensor([0, 0, 0, 0])