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

torch_geometric/utils/embedding.py

@@ -1,9 +1,11 @@
 import warnings
-from typing import Any, List
+from typing import Any, Dict, List, Optional, Type
 
 import torch
 from torch import Tensor
 
+from torch_geometric.typing import NodeType
+
 
 def get_embeddings(
     model: torch.nn.Module,
@@ -40,7 +42,8 @@ def get_embeddings(
             hook_handles.append(module.register_forward_hook(hook))
 
     if len(hook_handles) == 0:
-        warnings.warn("The 'model' does not have any 'MessagePassing' layers")
+        warnings.warn("The 'model' does not have any 'MessagePassing' layers",
+                      stacklevel=2)
 
     training = model.training
     model.eval()
@@ -52,3 +55,89 @@ def get_embeddings(
         handle.remove()
 
     return embeddings
+
+
+def get_embeddings_hetero(
+    model: torch.nn.Module,
+    supported_models: Optional[List[Type[torch.nn.Module]]] = None,
+    *args: Any,
+    **kwargs: Any,
+) -> Dict[NodeType, List[Tensor]]:
+    """Returns the output embeddings of all
+    :class:`~torch_geometric.nn.conv.MessagePassing` layers in a heterogeneous
+    :obj:`model`, organized by edge type.
+
+    Internally, this method registers forward hooks on all modules that process
+    heterogeneous graphs in the model and runs the forward pass of the model.
+    For heterogeneous models, the output is a dictionary where each key is a
+    node type and each value is a list of embeddings from different layers.
+
+    Args:
+        model (torch.nn.Module): The heterogeneous GNN model.
+        supported_models (List[Type[torch.nn.Module]], optional): A list of
+            supported model classes. If not provided, defaults to
+            [HGTConv, HANConv, HeteroConv].
+        *args: Arguments passed to the model.
+        **kwargs (optional): Additional keyword arguments passed to the model.
+
+    Returns:
+        Dict[NodeType, List[Tensor]]: A dictionary mapping each node type to
+        a list of embeddings from different layers.
+    """
+    from torch_geometric.nn import HANConv, HeteroConv, HGTConv
+    if not supported_models:
+        supported_models = [HGTConv, HANConv, HeteroConv]
+
+    # Dictionary to store node embeddings by type
+    node_embeddings_dict: Dict[NodeType, List[Tensor]] = {}
+
+    # Hook function to capture node embeddings
+    def hook(model: torch.nn.Module, inputs: Any, outputs: Any) -> None:
+        # Check if the outputs is a dictionary mapping node types to embeddings
+        if isinstance(outputs, dict) and outputs:
+            # Store embeddings for each node type
+            for node_type, embedding in outputs.items():
+                # Made sure that the outputs are a dictionary mapping node
+                # types to embeddings and remove the false positives.
+                if node_type not in node_embeddings_dict:
+                    node_embeddings_dict[node_type] = []
+                node_embeddings_dict[node_type].append(embedding.clone())
+
+    # List to store hook handles
+    hook_handles = []
+
+    # Find ModuleDict objects in the model
+    for _, module in model.named_modules():
+        # Handle the native heterogenous models, e.g. HGTConv, HANConv
+        # and HeteroConv, etc.
+        if isinstance(module, tuple(supported_models)):
+            hook_handles.append(module.register_forward_hook(hook))
+        else:
+            # Handle the heterogenous models that are generated by calling
+            # to_hetero() on the homogeneous models.
+            submodules = list(module.children())
+            submodules_contains_module_dict = any([
+                isinstance(submodule, torch.nn.ModuleDict)
+                for submodule in submodules
+            ])
+            if submodules_contains_module_dict:
+                hook_handles.append(module.register_forward_hook(hook))
+
+    if len(hook_handles) == 0:
+        warnings.warn(
+            "The 'model' does not have any heterogenous "
+            "'MessagePassing' layers", stacklevel=2)
+
+    # Run the model forward pass
+    training = model.training
+    model.eval()
+
+    with torch.no_grad():
+        model(*args, **kwargs)
+    model.train(training)
+
+    # Clean up hooks
+    for handle in hook_handles:
+        handle.remove()
+
+    return node_embeddings_dict

torch_geometric/utils/geodesic.py

@@ -2,6 +2,7 @@ import multiprocessing as mp
 import warnings
 from typing import Optional
 
+import numpy as np
 import torch
 from torch import Tensor
 
@@ -65,9 +66,10 @@ def geodesic_distance(  # noqa: D417
 
     if 'dest' in kwargs:
         dst = kwargs['dest']
-        warnings.warn("'dest' attribute in 'geodesic_distance' is deprecated "
-                      "and will be removed in a future release. Use the 'dst' "
-                      "argument instead.")
+        warnings.warn(
+            "'dest' attribute in 'geodesic_distance' is deprecated "
+            "and will be removed in a future release. Use the 'dst' "
+            "argument instead.", stacklevel=2)
 
     max_distance = float('inf') if max_distance is None else max_distance
 
@@ -82,54 +84,55 @@ def geodesic_distance(  # noqa: D417
 
     dtype = pos.dtype
 
-    pos = pos.detach().cpu().to(torch.double).numpy()
-    face = face.detach().t().cpu().to(torch.int).numpy()
+    pos_np = pos.detach().cpu().to(torch.double).numpy()
+    face_np = face.detach().t().cpu().to(torch.int).numpy()
 
     if src is None and dst is None:
-        out = gdist.local_gdist_matrix(pos, face,
-                                       max_distance * scale).toarray() / scale
+        out = gdist.local_gdist_matrix(
+            pos_np,
+            face_np,
+            max_distance * scale,
+        ).toarray() / scale
         return torch.from_numpy(out).to(dtype)
 
     if src is None:
-        src = torch.arange(pos.shape[0], dtype=torch.int).numpy()
+        src_np = torch.arange(pos.size(0), dtype=torch.int).numpy()
     else:
-        src = src.detach().cpu().to(torch.int).numpy()
-    assert src is not None
+        src_np = src.detach().cpu().to(torch.int).numpy()
 
-    dst = None if dst is None else dst.detach().cpu().to(torch.int).numpy()
+    dst_np = None if dst is None else dst.detach().cpu().to(torch.int).numpy()
 
     def _parallel_loop(
-        pos: Tensor,
-        face: Tensor,
-        src: Tensor,
-        dst: Optional[Tensor],
+        pos_np: np.ndarray,
+        face_np: np.ndarray,
+        src_np: np.ndarray,
+        dst_np: Optional[np.ndarray],
         max_distance: float,
         scale: float,
         i: int,
         dtype: torch.dtype,
     ) -> Tensor:
-        s = src[i:i + 1]
-        d = None if dst is None else dst[i:i + 1]
-        out = gdist.compute_gdist(pos, face, s, d, max_distance * scale)
+        s = src_np[i:i + 1]
+        d = None if dst_np is None else dst_np[i:i + 1]
+        out = gdist.compute_gdist(pos_np, face_np, s, d, max_distance * scale)
         out = out / scale
         return torch.from_numpy(out).to(dtype)
 
     num_workers = mp.cpu_count() if num_workers <= -1 else num_workers
     if num_workers > 0:
         with mp.Pool(num_workers) as pool:
-            outs = pool.starmap(
-                _parallel_loop,
-                [(pos, face, src, dst, max_distance, scale, i, dtype)
-                 for i in range(len(src))])
+            data = [(pos_np, face_np, src_np, dst_np, max_distance, scale, i,
+                     dtype) for i in range(len(src_np))]
+            outs = pool.starmap(_parallel_loop, data)
     else:
         outs = [
-            _parallel_loop(pos, face, src, dst, max_distance, scale, i, dtype)
-            for i in range(len(src))
+            _parallel_loop(pos_np, face_np, src_np, dst_np, max_distance,
+                           scale, i, dtype) for i in range(len(src_np))
         ]
 
     out = torch.cat(outs, dim=0)
 
     if dst is None:
-        out = out.view(-1, pos.shape[0])
+        out = out.view(-1, pos.size(0))
 
     return out

torch_geometric/utils/influence.py

@@ -0,0 +1,279 @@
+from typing import List, Tuple, Union, cast
+
+import torch
+from torch import Tensor
+from torch.autograd.functional import jacobian
+from tqdm.auto import tqdm
+
+from torch_geometric.data import Data
+from torch_geometric.utils import k_hop_subgraph
+
+
+def k_hop_subsets_rough(
+    node_idx: int,
+    num_hops: int,
+    edge_index: Tensor,
+    num_nodes: int,
+) -> List[Tensor]:
+    r"""Return *rough* (possibly overlapping) *k*-hop node subsets.
+
+    This is a thin wrapper around
+    :pyfunc:`torch_geometric.utils.k_hop_subgraph` that *additionally* returns
+    **all** intermediate hop subsets rather than the full union only.
+
+    Parameters
+    ----------
+    node_idx: int
+        Index or indices of the central node(s).
+    num_hops: int
+        Number of hops *k*.
+    edge_index: Tensor
+        Edge index in COO format with shape :math:`[2, \text{num_edges}]`.
+    num_nodes: int
+        Total number of nodes in the graph. Required to allocate the masks.
+
+    Returns:
+    -------
+    List[Tensor]
+        A list ``[H₀, H₁, …, H_k]`` where ``H₀`` contains the seed node(s) and
+        ``H_i`` (for *i*>0) contains **all** nodes that are exactly *i* hops
+        away in the *expanded* neighbourhood (i.e. overlaps are *not*
+        removed).
+    """
+    col, row = edge_index
+
+    node_mask = row.new_empty(num_nodes, dtype=torch.bool)
+    edge_mask = row.new_empty(row.size(0), dtype=torch.bool)
+
+    node_idx_ = torch.tensor([node_idx], device=row.device)
+
+    subsets = [node_idx_]
+    for _ in range(num_hops):
+        node_mask.zero_()
+        node_mask[subsets[-1]] = True
+        torch.index_select(node_mask, 0, row, out=edge_mask)
+        subsets.append(col[edge_mask])
+
+    return subsets
+
+
+def k_hop_subsets_exact(
+    node_idx: int,
+    num_hops: int,
+    edge_index: Tensor,
+    num_nodes: int,
+    device: Union[torch.device, str],
+) -> List[Tensor]:
+    """Return **disjoint** *k*-hop subsets.
+
+    This function refines :pyfunc:`k_hop_subsets_rough` by removing nodes that
+    have already appeared in previous hops, ensuring that each subset contains
+    nodes *exactly* *i* hops away from the seed.
+    """
+    rough_subsets = k_hop_subsets_rough(node_idx, num_hops, edge_index,
+                                        num_nodes)
+
+    exact_subsets: List[List[int]] = [rough_subsets[0].tolist()]
+    visited: set[int] = set(exact_subsets[0])
+
+    for hop_subset in rough_subsets[1:]:
+        fresh = set(hop_subset.tolist()) - visited
+        visited |= fresh
+        exact_subsets.append(list(fresh))
+
+    return [
+        torch.tensor(s, device=device, dtype=edge_index.dtype)
+        for s in exact_subsets
+    ]
+
+
+def jacobian_l1(
+    model: torch.nn.Module,
+    data: Data,
+    max_hops: int,
+    node_idx: int,
+    device: Union[torch.device, str],
+    *,
+    vectorize: bool = True,
+) -> Tensor:
+    """Compute the **L1 norm** of the Jacobian for a given node.
+
+    The Jacobian is evaluated w.r.t. the node features of the *k*-hop induced
+    sub‑graph centred at ``node_idx``. The result is *folded back* onto the
+    **original** node index space so that the returned tensor has length
+    ``data.num_nodes``, where the influence score will be zero for nodes
+    outside the *k*-hop subgraph.
+
+    Notes:
+    -----
+    *   The function assumes that the model *and* ``data.x`` share the same
+        floating‑point precision (e.g. both ``float32`` or both ``float16``).
+
+    """
+    # Build the induced *k*-hop sub‑graph (with node re‑labelling).
+    edge_index = cast(Tensor, data.edge_index)
+    x = cast(Tensor, data.x)
+    k_hop_nodes, sub_edge_index, mapping, _ = k_hop_subgraph(
+        node_idx, max_hops, edge_index, relabel_nodes=True)
+    # get the location of the *center* node inside the sub‑graph
+    root_pos = cast(int, mapping[0])
+
+    # Move tensors & model to the correct device
+    device = torch.device(device)
+    sub_x = x[k_hop_nodes].to(device)
+    sub_edge_index = sub_edge_index.to(device)
+    model = model.to(device)
+
+    # Jacobian evaluation
+    def _forward(x: Tensor) -> Tensor:
+        return model(x, sub_edge_index)[root_pos]
+
+    jac = jacobian(_forward, sub_x, vectorize=vectorize)
+    influence_sub = jac.abs().sum(dim=(0, 2))  # Sum of L1 norm
+    num_nodes = cast(int, data.num_nodes)
+    # Scatter the influence scores back to the *global* node space
+    influence_full = torch.zeros(num_nodes, dtype=influence_sub.dtype,
+                                 device=device)
+    influence_full[k_hop_nodes] = influence_sub
+
+    return influence_full
+
+
+def jacobian_l1_agg_per_hop(
+    model: torch.nn.Module,
+    data: Data,
+    max_hops: int,
+    node_idx: int,
+    device: Union[torch.device, str],
+    vectorize: bool = True,
+) -> Tensor:
+    """Aggregate Jacobian L1 norms **per hop** for node_idx.
+
+    Returns a vector ``[I_0, I_1, …, I_k]`` where ``I_i`` is the *total*
+    influence exerted by nodes that are exactly *i* hops away from
+    ``node_idx``.
+    """
+    num_nodes = cast(int, data.num_nodes)
+    edge_index = cast(Tensor, data.edge_index)
+    influence = jacobian_l1(model, data, max_hops, node_idx, device,
+                            vectorize=vectorize)
+    hop_subsets = k_hop_subsets_exact(node_idx, max_hops, edge_index,
+                                      num_nodes, influence.device)
+    single_node_influence_per_hop = [influence[s].sum() for s in hop_subsets]
+    return torch.tensor(single_node_influence_per_hop, device=influence.device)
+
+
+def avg_total_influence(
+    influence_all_nodes: Tensor,
+    normalize: bool = True,
+) -> Tensor:
+    """Compute the *influence‑weighted receptive field* ``R``."""
+    avg_total_influences = torch.mean(influence_all_nodes, dim=0)
+    if normalize:  # normalize by hop_0 (jacobian of the center node feature)
+        avg_total_influences = avg_total_influences / avg_total_influences[0]
+    return avg_total_influences
+
+
+def influence_weighted_receptive_field(T: Tensor) -> float:
+    """Compute the *influence‑weighted receptive field* ``R``.
+
+    Given an influence matrix ``T`` of shape ``[N, k+1]`` (i‑th row contains
+    the per‑hop influences of node *i*), the receptive field breadth *R* is
+    defined as the expected hop distance when weighting by influence.
+
+    A larger *R* indicates that, on average, influence comes from **farther**
+    hops.
+    """
+    normalised = T / torch.sum(T, dim=1, keepdim=True)
+    hops = torch.arange(T.shape[1]).float()  # 0 … k
+    breadth = normalised @ hops  # shape (N,)
+    return breadth.mean().item()
+
+
+def total_influence(
+    model: torch.nn.Module,
+    data: Data,
+    max_hops: int,
+    num_samples: Union[int, None] = None,
+    normalize: bool = True,
+    average: bool = True,
+    device: Union[torch.device, str] = "cpu",
+    vectorize: bool = True,
+) -> Tuple[Tensor, float]:
+    r"""Compute Jacobian‑based influence aggregates for *multiple* seed nodes,
+    as introduced in the
+    `"Towards Quantifying Long-Range Interactions in Graph Machine Learning:
+    a Large Graph Dataset and a Measurement"
+    <https://arxiv.org/abs/2503.09008>`_ paper.
+    This measurement quantifies how a GNN model's output at a node is
+    influenced by features of other nodes at increasing hop distances.
+
+    Specifically, for every sampled node :math:`v`, this method
+
+    1. evaluates the **L1‑norm** of the Jacobian of the model output at
+       :math:`v` w.r.t. the node features of its *k*-hop induced sub‑graph;
+    2. sums these scores **per hop** to obtain the influence vector
+       :math:`(I_{0}, I_{1}, \dots, I_{k})`;
+    3. optionally averages those vectors over all sampled nodes and
+       optionally normalises them by :math:`I_{0}`.
+
+    Please refer to Section 4 of the paper for a more detailed definition.
+
+    Args:
+        model (torch.nn.Module): A PyTorch Geometric‑compatible model with
+            forward signature ``model(x, edge_index) -> Tensor``.
+        data (torch_geometric.data.Data): Graph data object providing at least
+            :obj:`x` (node features) and :obj:`edge_index` (connectivity).
+        max_hops (int): Maximum hop distance :math:`k`.
+        num_samples (int, optional): Number of random seed nodes to evaluate.
+            If :obj:`None`, all nodes are used. (default: :obj:`None`)
+        normalize (bool, optional): If :obj:`True`, normalize each hop‑wise
+            influence by the influence of hop 0. (default: :obj:`True`)
+        average (bool, optional): If :obj:`True`, return the hop‑wise **mean**
+            over all seed nodes (shape ``[k+1]``).
+            If :obj:`False`, return the full influence matrix of shape
+            ``[N, k+1]``. (default: :obj:`True`)
+        device (torch.device or str, optional): Device on which to perform the
+            computation. (default: :obj:`"cpu"`)
+        vectorize (bool, optional): Forwarded to
+            :func:`torch.autograd.functional.jacobian`.  Keeping this
+            :obj:`True` is often faster but increases memory usage.
+            (default: :obj:`True`)
+
+    Returns:
+        Tuple[Tensor, float]:
+            * **avg_influence** (*Tensor*):
+              shape ``[k+1]`` if :obj:`average=True`;
+              shape ``[N, k+1]`` otherwise.
+            * **R** (*float*): Influence‑weighted receptive‑field breadth
+              returned by :func:`influence_weighted_receptive_field`.
+
+    Example::
+        >>> avg_I, R = total_influence(model, data, max_hops=3,
+        ...                            num_samples=1000)
+        >>> avg_I
+        tensor([1.0000, 0.1273, 0.0142, 0.0019])
+        >>> R
+        0.216
+    """
+    num_samples = data.num_nodes if num_samples is None else num_samples
+    num_nodes = cast(int, data.num_nodes)
+    nodes = torch.randperm(num_nodes)[:num_samples].tolist()
+
+    influence_all_nodes: List[Tensor] = [
+        jacobian_l1_agg_per_hop(model, data, max_hops, n, device,
+                                vectorize=vectorize)
+        for n in tqdm(nodes, desc="Influence")
+    ]
+    allnodes = torch.vstack(influence_all_nodes).detach().cpu()
+
+    # Average total influence at each hop
+    if average:
+        avg_influence = avg_total_influence(allnodes, normalize=normalize)
+    else:
+        avg_influence = allnodes
+
+    # Influence‑weighted receptive field
+    R = influence_weighted_receptive_field(allnodes)
+
+    return avg_influence, R

torch_geometric/utils/map.py

@@ -1,4 +1,3 @@
-import warnings
 from typing import Optional, Tuple, Union
 
 import numpy as np
@@ -6,6 +5,10 @@ import torch
 from torch import Tensor
 from torch.utils.dlpack import from_dlpack
 
+from torch_geometric.warnings import WarningCache
+
+_warning_cache = WarningCache()
+
 
 def map_index(
     src: Tensor,
@@ -14,7 +17,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:
@@ -93,10 +96,10 @@ def map_index(
             WITH_CUDF = True
         except ImportError:
             import pandas as pd
-            warnings.warn("Using CPU-based processing within 'map_index' "
-                          "which may cause slowdowns and device "
-                          "synchronization. Consider installing 'cudf' to "
-                          "accelerate computation")
+            _warning_cache.warn("Using CPU-based processing within "
+                                "'map_index' which may cause slowdowns and "
+                                "device synchronization. Consider installing "
+                                "'cudf' to accelerate computation")
     else:
         import pandas as pd
 
@@ -148,10 +151,11 @@ def map_index(
         if inclusive:
             try:
                 out = from_dlpack(result['right_ser'].to_dlpack())
-            except ValueError:
-                raise ValueError("Found invalid entries in 'src' that do not "
-                                 "have a corresponding entry in 'index'. Set "
-                                 "`inclusive=False` to ignore these entries.")
+            except ValueError as e:
+                raise ValueError(
+                    "Found invalid entries in 'src' that do not "
+                    "have a corresponding entry in 'index'. Set "
+                    "`inclusive=False` to ignore these entries.") from e
         else:
             out = from_dlpack(result['right_ser'].fillna(-1).to_dlpack())
 

torch_geometric/utils/nested.py

@@ -43,7 +43,7 @@ def to_nested_tensor(
         xs = [x]
 
     # This currently copies the data, although `x` is already contiguous.
-    # Sadly, there does not exist any (public) API to preven this :(
+    # Sadly, there does not exist any (public) API to prevent this :(
     return torch.nested.as_nested_tensor(xs)
 
 

torch_geometric/utils/smiles.py

@@ -91,7 +91,7 @@ def from_rdmol(mol: Any) -> 'torch_geometric.data.Data':
     assert isinstance(mol, Chem.Mol)
 
     xs: List[List[int]] = []
-    for atom in mol.GetAtoms():  # type: ignore
+    for atom in mol.GetAtoms():
         row: List[int] = []
         row.append(x_map['atomic_num'].index(atom.GetAtomicNum()))
         row.append(x_map['chirality'].index(str(atom.GetChiralTag())))
@@ -108,7 +108,7 @@ def from_rdmol(mol: Any) -> 'torch_geometric.data.Data':
     x = torch.tensor(xs, dtype=torch.long).view(-1, 9)
 
     edge_indices, edge_attrs = [], []
-    for bond in mol.GetBonds():  # type: ignore
+    for bond in mol.GetBonds():
         i = bond.GetBeginAtomIdx()
         j = bond.GetEndAtomIdx()
 
@@ -148,7 +148,7 @@ def from_smiles(
     """
     from rdkit import Chem, RDLogger
 
-    RDLogger.DisableLog('rdApp.*')  # type: ignore
+    RDLogger.DisableLog('rdApp.*')  # type: ignore[attr-defined]
 
     mol = Chem.MolFromSmiles(smiles)