pyg-nightly 2.7.0.dev20241009__py3-none-any.whl → 2.8.0.dev20251207__py3-none-any.whl

torch_geometric/utils/__init__.py

@@ -53,10 +53,11 @@ from ._negative_sampling import (negative_sampling, batched_negative_sampling,
                                  structured_negative_sampling_feasible)
 from .augmentation import shuffle_node, mask_feature, add_random_edge
 from ._tree_decomposition import tree_decomposition
-from .embedding import get_embeddings
+from .embedding import get_embeddings, get_embeddings_hetero
 from ._trim_to_layer import trim_to_layer
 from .ppr import get_ppr
 from ._train_test_split_edges import train_test_split_edges
+from .influence import total_influence
 
 __all__ = [
     'scatter',
@@ -145,9 +146,11 @@ __all__ = [
     'add_random_edge',
     'tree_decomposition',
     'get_embeddings',
+    'get_embeddings_hetero',
     'trim_to_layer',
     'get_ppr',
     'train_test_split_edges',
+    'total_influence',
 ]
 
 # `structured_negative_sampling_feasible` is a long name and thus destroys the

torch_geometric/utils/_lexsort.py

@@ -1,11 +1,7 @@
 from typing import List
 
-import numpy as np
-import torch
 from torch import Tensor
 
-import torch_geometric.typing
-
 
 def lexsort(
     keys: List[Tensor],
@@ -28,11 +24,6 @@ def lexsort(
     """
     assert len(keys) >= 1
 
-    if not torch_geometric.typing.WITH_PT113:
-        keys = [k.neg() for k in keys] if descending else keys
-        out = np.lexsort([k.detach().cpu().numpy() for k in keys], axis=dim)
-        return torch.from_numpy(out).to(keys[0].device)
-
     out = keys[0].argsort(dim=dim, descending=descending, stable=True)
     for k in keys[1:]:
         index = k.gather(dim, out)

torch_geometric/utils/_negative_sampling.py

@@ -12,7 +12,7 @@ from torch_geometric.utils.num_nodes import maybe_num_nodes
 def negative_sampling(
     edge_index: Tensor,
     num_nodes: Optional[Union[int, Tuple[int, int]]] = None,
-    num_neg_samples: Optional[int] = None,
+    num_neg_samples: Optional[Union[int, float]] = None,
     method: str = "sparse",
     force_undirected: bool = False,
 ) -> Tensor:
@@ -25,10 +25,12 @@ def negative_sampling(
             If given as a tuple, then :obj:`edge_index` is interpreted as a
             bipartite graph with shape :obj:`(num_src_nodes, num_dst_nodes)`.
             (default: :obj:`None`)
-        num_neg_samples (int, optional): The (approximate) number of negative
-            samples to return.
-            If set to :obj:`None`, will try to return a negative edge for every
-            positive edge. (default: :obj:`None`)
+        num_neg_samples (int or float, optional): The (approximate) number of
+            negative samples to return. If set to a floating-point value, it
+            represents the ratio of negative samples to generate based on the
+            number of positive edges. If set to :obj:`None`, will try to
+            return a negative edge for every positive edge.
+            (default: :obj:`None`)
         method (str, optional): The method to use for negative sampling,
             *i.e.* :obj:`"sparse"` or :obj:`"dense"`.
             This is a memory/runtime trade-off.
@@ -48,6 +50,11 @@ def negative_sampling(
         tensor([[3, 0, 0, 3],
                 [2, 3, 2, 1]])
 
+        >>> negative_sampling(edge_index, num_nodes=(3, 4),
+        ...                   num_neg_samples=0.5)  # 50% of positive edges
+        tensor([[0, 3],
+                [3, 0]])
+
         >>> # For bipartite graph
         >>> negative_sampling(edge_index, num_nodes=(3, 4))
         tensor([[0, 2, 2, 1],
@@ -74,6 +81,8 @@ def negative_sampling(
 
     if num_neg_samples is None:
         num_neg_samples = edge_index.size(1)
+    elif isinstance(num_neg_samples, float):
+        num_neg_samples = int(num_neg_samples * edge_index.size(1))
     if force_undirected:
         num_neg_samples = num_neg_samples // 2
 
@@ -100,10 +109,9 @@ def negative_sampling(
         idx = idx.to('cpu')
         for _ in range(3):  # Number of tries to sample negative indices.
             rnd = sample(population, sample_size, device='cpu')
-            mask = np.isin(rnd.numpy(), idx.numpy())  # type: ignore
+            mask = torch.from_numpy(np.isin(rnd.numpy(), idx.numpy())).bool()
             if neg_idx is not None:
-                mask |= np.isin(rnd, neg_idx.to('cpu'))
-            mask = torch.from_numpy(mask).to(torch.bool)
+                mask |= torch.from_numpy(np.isin(rnd, neg_idx.cpu())).bool()
             rnd = rnd[~mask].to(edge_index.device)
             neg_idx = rnd if neg_idx is None else torch.cat([neg_idx, rnd])
             if neg_idx.numel() >= num_neg_samples:
@@ -117,7 +125,7 @@ def negative_sampling(
 def batched_negative_sampling(
     edge_index: Tensor,
     batch: Union[Tensor, Tuple[Tensor, Tensor]],
-    num_neg_samples: Optional[int] = None,
+    num_neg_samples: Optional[Union[int, float]] = None,
     method: str = "sparse",
     force_undirected: bool = False,
 ) -> Tensor:
@@ -131,9 +139,11 @@ def batched_negative_sampling(
             node to a specific example.
             If given as a tuple, then :obj:`edge_index` is interpreted as a
             bipartite graph connecting two different node types.
-        num_neg_samples (int, optional): The number of negative samples to
-            return. If set to :obj:`None`, will try to return a negative edge
-            for every positive edge. (default: :obj:`None`)
+        num_neg_samples (int or float, optional): The number of negative
+            samples to return. If set to :obj:`None`, will try to return a
+            negative edge for every positive edge. If float, it will generate
+            :obj:`num_neg_samples * num_edges` negative samples.
+            (default: :obj:`None`)
         method (str, optional): The method to use for negative sampling,
             *i.e.* :obj:`"sparse"` or :obj:`"dense"`.
             This is a memory/runtime trade-off.
@@ -157,6 +167,11 @@ def batched_negative_sampling(
         tensor([[3, 1, 3, 2, 7, 7, 6, 5],
                 [2, 0, 1, 1, 5, 6, 4, 4]])
 
+        >>> # Using float multiplier for negative samples
+        >>> batched_negative_sampling(edge_index, batch, num_neg_samples=1.5)
+        tensor([[3, 1, 3, 2, 7, 7, 6, 5, 2, 0, 1, 1],
+                [2, 0, 1, 1, 5, 6, 4, 4, 3, 2, 3, 0]])
+
         >>> # For bipartite graph
         >>> edge_index1 = torch.as_tensor([[0, 0, 1, 1], [0, 1, 2, 3]])
         >>> edge_index2 = edge_index1 + torch.tensor([[2], [4]])

torch_geometric/utils/_scatter.py

@@ -8,185 +8,134 @@ from torch_geometric import is_compiling, is_in_onnx_export, warnings
 from torch_geometric.typing import torch_scatter
 from torch_geometric.utils.functions import cumsum
 
-if torch_geometric.typing.WITH_PT112:  # pragma: no cover
-
-    warnings.filterwarnings('ignore', '.*is in beta and the API may change.*')
-
-    def scatter(
-        src: Tensor,
-        index: Tensor,
-        dim: int = 0,
-        dim_size: Optional[int] = None,
-        reduce: str = 'sum',
-    ) -> Tensor:
-        r"""Reduces all values from the :obj:`src` tensor at the indices
-        specified in the :obj:`index` tensor along a given dimension
-        :obj:`dim`. See the `documentation
-        <https://pytorch-scatter.readthedocs.io/en/latest/functions/
-        scatter.html>`__ of the :obj:`torch_scatter` package for more
-        information.
-
-        Args:
-            src (torch.Tensor): The source tensor.
-            index (torch.Tensor): The index tensor.
-            dim (int, optional): The dimension along which to index.
-                (default: :obj:`0`)
-            dim_size (int, optional): The size of the output tensor at
-                dimension :obj:`dim`. If set to :obj:`None`, will create a
-                minimal-sized output tensor according to
-                :obj:`index.max() + 1`. (default: :obj:`None`)
-            reduce (str, optional): The reduce operation (:obj:`"sum"`,
-                :obj:`"mean"`, :obj:`"mul"`, :obj:`"min"` or :obj:`"max"`,
-                :obj:`"any"`). (default: :obj:`"sum"`)
-        """
-        if isinstance(index, Tensor) and index.dim() != 1:
-            raise ValueError(f"The `index` argument must be one-dimensional "
-                             f"(got {index.dim()} dimensions)")
-
-        dim = src.dim() + dim if dim < 0 else dim
-
-        if isinstance(src, Tensor) and (dim < 0 or dim >= src.dim()):
-            raise ValueError(f"The `dim` argument must lay between 0 and "
-                             f"{src.dim() - 1} (got {dim})")
-
-        if dim_size is None:
-            dim_size = int(index.max()) + 1 if index.numel() > 0 else 0
-
-        # For now, we maintain various different code paths, based on whether
-        # the input requires gradients and whether it lays on the CPU/GPU.
-        # For example, `torch_scatter` is usually faster than
-        # `torch.scatter_reduce` on GPU, while `torch.scatter_reduce` is faster
-        # on CPU.
-        # `torch.scatter_reduce` has a faster forward implementation for
-        # "min"/"max" reductions since it does not compute additional arg
-        # indices, but is therefore way slower in its backward implementation.
-        # More insights can be found in `test/utils/test_scatter.py`.
-
-        size = src.size()[:dim] + (dim_size, ) + src.size()[dim + 1:]
-
-        # For "any" reduction, we use regular `scatter_`:
-        if reduce == 'any':
-            index = broadcast(index, src, dim)
-            return src.new_zeros(size).scatter_(dim, index, src)
+warnings.filterwarnings('ignore', '.*is in beta and the API may change.*')
 
-        # For "sum" and "mean" reduction, we make use of `scatter_add_`:
-        if reduce == 'sum' or reduce == 'add':
-            index = broadcast(index, src, dim)
-            return src.new_zeros(size).scatter_add_(dim, index, src)
 
-        if reduce == 'mean':
-            count = src.new_zeros(dim_size)
-            count.scatter_add_(0, index, src.new_ones(src.size(dim)))
-            count = count.clamp(min=1)
+def scatter(
+    src: Tensor,
+    index: Tensor,
+    dim: int = 0,
+    dim_size: Optional[int] = None,
+    reduce: str = 'sum',
+) -> Tensor:
+    r"""Reduces all values from the :obj:`src` tensor at the indices specified
+    in the :obj:`index` tensor along a given dimension ``dim``. See the
+    `documentation <https://pytorch-scatter.readthedocs.io/en/latest/functions/scatter.html>`__  # noqa: E501
+    of the ``torch_scatter`` package for more information.
 
-            index = broadcast(index, src, dim)
-            out = src.new_zeros(size).scatter_add_(dim, index, src)
-
-            return out / broadcast(count, out, dim)
-
-        # For "min" and "max" reduction, we prefer `scatter_reduce_` on CPU or
-        # in case the input does not require gradients:
-        if reduce in ['min', 'max', 'amin', 'amax']:
-            if (not torch_geometric.typing.WITH_TORCH_SCATTER
-                    or is_compiling() or is_in_onnx_export() or not src.is_cuda
-                    or not src.requires_grad):
-
-                if (src.is_cuda and src.requires_grad and not is_compiling()
-                        and not is_in_onnx_export()):
-                    warnings.warn(f"The usage of `scatter(reduce='{reduce}')` "
-                                  f"can be accelerated via the 'torch-scatter'"
-                                  f" package, but it was not found")
-
-                index = broadcast(index, src, dim)
-                if not is_in_onnx_export():
-                    return src.new_zeros(size).scatter_reduce_(
-                        dim, index, src, reduce=f'a{reduce[-3:]}',
-                        include_self=False)
-
-                fill = torch.full(  # type: ignore
-                    size=(1, ),
-                    fill_value=src.min() if 'max' in reduce else src.max(),
-                    dtype=src.dtype,
-                    device=src.device,
-                ).expand_as(src)
-                out = src.new_zeros(size).scatter_reduce_(
-                    dim, index, fill, reduce=f'a{reduce[-3:]}',
-                    include_self=True)
-                return out.scatter_reduce_(dim, index, src,
-                                           reduce=f'a{reduce[-3:]}',
-                                           include_self=True)
-
-            return torch_scatter.scatter(src, index, dim, dim_size=dim_size,
-                                         reduce=reduce[-3:])
-
-        # For "mul" reduction, we prefer `scatter_reduce_` on CPU:
-        if reduce == 'mul':
-            if (not torch_geometric.typing.WITH_TORCH_SCATTER
-                    or is_compiling() or not src.is_cuda):
-
-                if src.is_cuda and not is_compiling():
-                    warnings.warn(f"The usage of `scatter(reduce='{reduce}')` "
-                                  f"can be accelerated via the 'torch-scatter'"
-                                  f" package, but it was not found")
-
-                index = broadcast(index, src, dim)
-                # We initialize with `one` here to match `scatter_mul` output:
-                return src.new_ones(size).scatter_reduce_(
-                    dim, index, src, reduce='prod', include_self=True)
-
-            return torch_scatter.scatter(src, index, dim, dim_size=dim_size,
-                                         reduce='mul')
-
-        raise ValueError(f"Encountered invalid `reduce` argument '{reduce}'")
-
-else:  # pragma: no cover
-
-    def scatter(
-        src: Tensor,
-        index: Tensor,
-        dim: int = 0,
-        dim_size: Optional[int] = None,
-        reduce: str = 'sum',
-    ) -> Tensor:
-        r"""Reduces all values from the :obj:`src` tensor at the indices
-        specified in the :obj:`index` tensor along a given dimension
-        :obj:`dim`. See the `documentation
-        <https://pytorch-scatter.readthedocs.io/en/latest/functions/
-        scatter.html>`_ of the :obj:`torch_scatter` package for more
-        information.
-
-        Args:
-            src (torch.Tensor): The source tensor.
-            index (torch.Tensor): The index tensor.
-            dim (int, optional): The dimension along which to index.
-                (default: :obj:`0`)
-            dim_size (int, optional): The size of the output tensor at
-                dimension :obj:`dim`. If set to :obj:`None`, will create a
-                minimal-sized output tensor according to
-                :obj:`index.max() + 1`. (default: :obj:`None`)
-            reduce (str, optional): The reduce operation (:obj:`"sum"`,
-                :obj:`"mean"`, :obj:`"mul"`, :obj:`"min"` or :obj:`"max"`,
-                :obj:`"any"`). (default: :obj:`"sum"`)
-        """
-        if reduce == 'any':
-            dim = src.dim() + dim if dim < 0 else dim
-
-            if dim_size is None:
-                dim_size = int(index.max()) + 1 if index.numel() > 0 else 0
-
-            size = src.size()[:dim] + (dim_size, ) + src.size()[dim + 1:]
+    Args:
+        src (torch.Tensor): The source tensor.
+        index (torch.Tensor): The index tensor.
+        dim (int, optional): The dimension along which to index.
+            (default: ``0``)
+        dim_size (int, optional): The size of the output tensor at dimension
+            ``dim``. If set to :obj:`None`, will create a minimal-sized output
+            tensor according to ``index.max() + 1``. (default: :obj:`None`)
+        reduce (str, optional): The reduce operation (``"sum"``, ``"mean"``,
+            ``"mul"``, ``"min"``, ``"max"`` or ``"any"``). (default: ``"sum"``)
+    """
+    if isinstance(index, Tensor) and index.dim() != 1:
+        raise ValueError(f"The `index` argument must be one-dimensional "
+                         f"(got {index.dim()} dimensions)")
+
+    dim = src.dim() + dim if dim < 0 else dim
+
+    if isinstance(src, Tensor) and (dim < 0 or dim >= src.dim()):
+        raise ValueError(f"The `dim` argument must lay between 0 and "
+                         f"{src.dim() - 1} (got {dim})")
+
+    if dim_size is None:
+        dim_size = int(index.max()) + 1 if index.numel() > 0 else 0
+
+    # For now, we maintain various different code paths, based on whether
+    # the input requires gradients and whether it lays on the CPU/GPU.
+    # For example, `torch_scatter` is usually faster than
+    # `torch.scatter_reduce` on GPU, while `torch.scatter_reduce` is faster
+    # on CPU.
+    # `torch.scatter_reduce` has a faster forward implementation for
+    # "min"/"max" reductions since it does not compute additional arg
+    # indices, but is therefore way slower in its backward implementation.
+    # More insights can be found in `test/utils/test_scatter.py`.
+
+    size = src.size()[:dim] + (dim_size, ) + src.size()[dim + 1:]
+
+    # For "any" reduction, we use regular `scatter_`:
+    if reduce == 'any':
+        index = broadcast(index, src, dim)
+        return src.new_zeros(size).scatter_(dim, index, src)
+
+    # For "sum" and "mean" reduction, we make use of `scatter_add_`:
+    if reduce == 'sum' or reduce == 'add':
+        index = broadcast(index, src, dim)
+        return src.new_zeros(size).scatter_add_(dim, index, src)
+
+    if reduce == 'mean':
+        count = src.new_zeros(dim_size)
+        count.scatter_add_(0, index, src.new_ones(src.size(dim)))
+        count = count.clamp(min=1)
+
+        index = broadcast(index, src, dim)
+        out = src.new_zeros(size).scatter_add_(dim, index, src)
+
+        return out / broadcast(count, out, dim)
+
+    # For "min" and "max" reduction, we prefer `scatter_reduce_` on CPU or
+    # in case the input does not require gradients:
+    if reduce in ['min', 'max', 'amin', 'amax']:
+        if (not torch_geometric.typing.WITH_TORCH_SCATTER or is_compiling()
+                or is_in_onnx_export() or not src.is_cuda
+                or not src.requires_grad):
+
+            if (src.is_cuda and src.requires_grad and not is_compiling()
+                    and not is_in_onnx_export()):
+                warnings.warn(
+                    f"The usage of `scatter(reduce='{reduce}')` "
+                    f"can be accelerated via the 'torch-scatter'"
+                    f" package, but it was not found", stacklevel=2)
 
             index = broadcast(index, src, dim)
-            return src.new_zeros(size).scatter_(dim, index, src)
+            if not is_in_onnx_export():
+                return src.new_zeros(size).scatter_reduce_(
+                    dim, index, src, reduce=f'a{reduce[-3:]}',
+                    include_self=False)
+
+            fill = torch.full(  # type: ignore
+                size=(1, ),
+                fill_value=src.min() if 'max' in reduce else src.max(),
+                dtype=src.dtype,
+                device=src.device,
+            ).expand_as(src)
+            out = src.new_zeros(size).scatter_reduce_(dim, index, fill,
+                                                      reduce=f'a{reduce[-3:]}',
+                                                      include_self=True)
+            return out.scatter_reduce_(dim, index, src,
+                                       reduce=f'a{reduce[-3:]}',
+                                       include_self=True)
 
-        if not torch_geometric.typing.WITH_TORCH_SCATTER:
-            raise ImportError("'scatter' requires the 'torch-scatter' package")
+        return torch_scatter.scatter(src, index, dim, dim_size=dim_size,
+                                     reduce=reduce[-3:])
+
+    # For "mul" reduction, we prefer `scatter_reduce_` on CPU:
+    if reduce == 'mul':
+        if (not torch_geometric.typing.WITH_TORCH_SCATTER or is_compiling()
+                or not src.is_cuda):
+
+            if src.is_cuda and not is_compiling():
+                warnings.warn(
+                    f"The usage of `scatter(reduce='{reduce}')` "
+                    f"can be accelerated via the 'torch-scatter'"
+                    f" package, but it was not found", stacklevel=2)
 
-        if reduce == 'amin' or reduce == 'amax':
-            reduce = reduce[-3:]
+            index = broadcast(index, src, dim)
+            # We initialize with `one` here to match `scatter_mul` output:
+            return src.new_ones(size).scatter_reduce_(dim, index, src,
+                                                      reduce='prod',
+                                                      include_self=True)
 
         return torch_scatter.scatter(src, index, dim, dim_size=dim_size,
-                                     reduce=reduce)
+                                     reduce='mul')
+
+    raise ValueError(f"Encountered invalid `reduce` argument '{reduce}'")
 
 
 def broadcast(src: Tensor, ref: Tensor, dim: int) -> Tensor:
@@ -215,24 +164,18 @@ def scatter_argmax(
     if dim_size is None:
         dim_size = int(index.max()) + 1 if index.numel() > 0 else 0
 
-    if torch_geometric.typing.WITH_PT112:
-        if not is_in_onnx_export():
-            res = src.new_empty(dim_size)
-            res.scatter_reduce_(0, index, src.detach(), reduce='amax',
-                                include_self=False)
-        else:
-            # `include_self=False` is currently not supported by ONNX:
-            res = src.new_full(
-                size=(dim_size, ),
-                fill_value=src.min(),  # type: ignore
-            )
-            res.scatter_reduce_(0, index, src.detach(), reduce="amax",
-                                include_self=True)
-    elif torch_geometric.typing.WITH_PT111:
-        res = torch.scatter_reduce(src.detach(), 0, index, reduce='amax',
-                                   output_size=dim_size)  # type: ignore
+    if not is_in_onnx_export():
+        res = src.new_empty(dim_size)
+        res.scatter_reduce_(0, index, src.detach(), reduce='amax',
+                            include_self=False)
     else:
-        raise ValueError("'scatter_argmax' requires PyTorch >= 1.11")
+        # `include_self=False` is currently not supported by ONNX:
+        res = src.new_full(
+            size=(dim_size, ),
+            fill_value=src.min(),  # type: ignore
+        )
+        res.scatter_reduce_(0, index, src.detach(), reduce="amax",
+                            include_self=True)
 
     out = index.new_full((dim_size, ), fill_value=dim_size - 1)
     nonzero = (src == res[index]).nonzero().view(-1)
@@ -290,13 +233,7 @@ def group_argsort(
 
     # Compute `grouped_argsort`:
     src = src - 2 * index if descending else src + 2 * index
-    if torch_geometric.typing.WITH_PT113:
-        perm = src.argsort(descending=descending, stable=stable)
-    else:
-        perm = src.argsort(descending=descending)
-        if stable:
-            warnings.warn("Ignoring option `stable=True` in 'group_argsort' "
-                          "since it requires PyTorch >= 1.13.0")
+    perm = src.argsort(descending=descending, stable=stable)
     out = torch.empty_like(index)
     out[perm] = torch.arange(index.numel(), device=index.device)
 
@@ -351,5 +288,5 @@ def group_cat(
     """
     assert len(tensors) == len(indices)
     index, perm = torch.cat(indices).sort(stable=True)
-    out = torch.cat(tensors, dim=0)[perm]
+    out = torch.cat(tensors, dim=dim).index_select(dim, perm)
     return (out, index) if return_index else out

torch_geometric/utils/_sort_edge_index.py

@@ -107,8 +107,6 @@ def sort_edge_index(  # noqa: F811
     num_nodes = maybe_num_nodes(edge_index, num_nodes)
 
     if num_nodes * num_nodes > torch_geometric.typing.MAX_INT64:
-        if not torch_geometric.typing.WITH_PT113:
-            raise ValueError("'sort_edge_index' will result in an overflow")
         perm = lexsort(keys=[
             edge_index[int(sort_by_row)],
             edge_index[1 - int(sort_by_row)],

torch_geometric/utils/_spmm.py

@@ -63,18 +63,20 @@ def spmm(
 
         # Always convert COO to CSR for more efficient processing:
         if src.layout == torch.sparse_coo:
-            warnings.warn(f"Converting sparse tensor to CSR format for more "
-                          f"efficient processing. Consider converting your "
-                          f"sparse tensor to CSR format beforehand to avoid "
-                          f"repeated conversion (got '{src.layout}')")
+            warnings.warn(
+                f"Converting sparse tensor to CSR format for more "
+                f"efficient processing. Consider converting your "
+                f"sparse tensor to CSR format beforehand to avoid "
+                f"repeated conversion (got '{src.layout}')", stacklevel=2)
             src = src.to_sparse_csr()
 
         # Warn in case of CSC format without gradient computation:
         if src.layout == torch.sparse_csc and not other.requires_grad:
-            warnings.warn(f"Converting sparse tensor to CSR format for more "
-                          f"efficient processing. Consider converting your "
-                          f"sparse tensor to CSR format beforehand to avoid "
-                          f"repeated conversion (got '{src.layout}')")
+            warnings.warn(
+                f"Converting sparse tensor to CSR format for more "
+                f"efficient processing. Consider converting your "
+                f"sparse tensor to CSR format beforehand to avoid "
+                f"repeated conversion (got '{src.layout}')", stacklevel=2)
 
         # Use the default code path for `sum` reduction (works on CPU/GPU):
         if reduce == 'sum':
@@ -99,10 +101,11 @@ def spmm(
         # TODO The `torch.sparse.mm` code path with the `reduce` argument does
         # not yet support CSC :(
         if src.layout == torch.sparse_csc:
-            warnings.warn(f"Converting sparse tensor to CSR format for more "
-                          f"efficient processing. Consider converting your "
-                          f"sparse tensor to CSR format beforehand to avoid "
-                          f"repeated conversion (got '{src.layout}')")
+            warnings.warn(
+                f"Converting sparse tensor to CSR format for more "
+                f"efficient processing. Consider converting your "
+                f"sparse tensor to CSR format beforehand to avoid "
+                f"repeated conversion (got '{src.layout}')", stacklevel=2)
             src = src.to_sparse_csr()
 
         return torch.sparse.mm(src, other, reduce)
@@ -115,8 +118,7 @@ def spmm(
         if src.layout == torch.sparse_csr:
             ptr = src.crow_indices()
             deg = ptr[1:] - ptr[:-1]
-        elif (torch_geometric.typing.WITH_PT112
-              and src.layout == torch.sparse_csc):
+        elif src.layout == torch.sparse_csc:
             assert src.layout == torch.sparse_csc
             ones = torch.ones_like(src.values())
             index = src.row_indices()

torch_geometric/utils/_subgraph.py

@@ -346,10 +346,12 @@ def k_hop_subgraph(
 
     subsets = [node_idx]
 
+    preserved_edge_mask = torch.zeros_like(edge_mask)
     for _ in range(num_hops):
         node_mask.fill_(False)
         node_mask[subsets[-1]] = True
         torch.index_select(node_mask, 0, row, out=edge_mask)
+        preserved_edge_mask |= edge_mask
         subsets.append(col[edge_mask])
 
     subset, inv = torch.cat(subsets).unique(return_inverse=True)
@@ -360,6 +362,8 @@ def k_hop_subgraph(
 
     if not directed:
         edge_mask = node_mask[row] & node_mask[col]
+    else:
+        edge_mask = preserved_edge_mask
 
     edge_index = edge_index[:, edge_mask]
 

torch_geometric/utils/_trim_to_layer.py

@@ -234,10 +234,10 @@ def trim_sparse_tensor(src: SparseTensor, size: Tuple[int, int],
     rowptr = torch.narrow(rowptr, 0, 0, size[0] + 1).clone()
     rowptr[num_seed_nodes + 1:] = rowptr[num_seed_nodes]
 
-    col = torch.narrow(col, 0, 0, rowptr[-1])
+    col = torch.narrow(col, 0, 0, rowptr[-1])  # type: ignore
 
     if value is not None:
-        value = torch.narrow(value, 0, 0, rowptr[-1])
+        value = torch.narrow(value, 0, 0, rowptr[-1])  # type: ignore
 
     csr2csc = src.storage._csr2csc
     if csr2csc is not None: