code: def k_hop_subgraph(
      node_idx: Union[int, List[int], Tensor],
      num_hops: int,
      edge_index: Tensor,
      relabel_nodes: bool = False,
      num_nodes: Optional[int] = None,
      flow: str = "source_to_target",
      directed: bool = False,
  ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
      r"""Computes the induced subgraph of :obj:`edge_index` around all nodes in
      :attr:`node_idx` reachable within :math:`k` hops.
  
      The :attr:`flow` argument denotes the direction of edges for finding
      :math:`k`-hop neighbors. If set to :obj:`"source_to_target"`, then the
      method will find all neighbors that point to the initial set of seed nodes
      in :attr:`node_idx.`
      This mimics the natural flow of message passing in Graph Neural Networks.
  
      The method returns (1) the nodes involved in the subgraph, (2) the filtered
      :obj:`edge_index` connectivity, (3) the mapping from node indices in
      :obj:`node_idx` to their new location, and (4) the edge mask indicating
      which edges were preserved.
  
      Args:
          node_idx (int, list, tuple or :obj:`torch.Tensor`): The central seed
              node(s).
          num_hops (int): The number of hops :math:`k`.
          edge_index (LongTensor): The edge indices.
          relabel_nodes (bool, optional): If set to :obj:`True`, the resulting
              :obj:`edge_index` will be relabeled to hold consecutive indices
              starting from zero. (default: :obj:`False`)
          num_nodes (int, optional): The number of nodes, *i.e.*
              :obj:`max_val + 1` of :attr:`edge_index`. (default: :obj:`None`)
          flow (str, optional): The flow direction of :math:`k`-hop aggregation
              (:obj:`"source_to_target"` or :obj:`"target_to_source"`).
              (default: :obj:`"source_to_target"`)
          directed (bool, optional): If set to :obj:`True`, will only include
              directed edges to the seed nodes :obj:`node_idx`.
              (default: :obj:`False`)
  
      :rtype: (:class:`LongTensor`, :class:`LongTensor`, :class:`LongTensor`,
               :class:`BoolTensor`)
  
      Examples:
          >>> edge_index = torch.tensor([[0, 1, 2, 3, 4, 5],
          ...                            [2, 2, 4, 4, 6, 6]])
  
          >>> # Center node 6, 2-hops
          >>> subset, edge_index, mapping, edge_mask = k_hop_subgraph(
          ...     6, 2, edge_index, relabel_nodes=True)
          >>> subset
          tensor([2, 3, 4, 5, 6])
          >>> edge_index
          tensor([[0, 1, 2, 3],
                  [2, 2, 4, 4]])
          >>> mapping
          tensor([4])
          >>> edge_mask
          tensor([False, False,  True,  True,  True,  True])
          >>> subset[mapping]
          tensor([6])
  
          >>> edge_index = torch.tensor([[1, 2, 4, 5],
          ...                            [0, 1, 5, 6]])
          >>> (subset, edge_index,
          ...  mapping, edge_mask) = k_hop_subgraph([0, 6], 2,
          ...                                       edge_index,
          ...                                       relabel_nodes=True)
          >>> subset
          tensor([0, 1, 2, 4, 5, 6])
          >>> edge_index
          tensor([[1, 2, 3, 4],
                  [0, 1, 4, 5]])
          >>> mapping
          tensor([0, 5])
          >>> edge_mask
          tensor([True, True, True, True])
          >>> subset[mapping]
          tensor([0, 6])
      """
      num_nodes = maybe_num_nodes(edge_index, num_nodes)
  
      assert flow in ["source_to_target", "target_to_source"]
      if flow == "target_to_source":
          row, col = edge_index
      else:
          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)
  
      if isinstance(node_idx, int):
          node_idx = torch.tensor([node_idx], device=row.device)
      elif isinstance(node_idx, (list, tuple)):
          node_idx = torch.tensor(node_idx, device=row.device)
      else:
          node_idx = node_idx.to(row.device)
  
      subsets = [node_idx]
  
      for _ in range(num_hops):
          node_mask.fill_(False)
          node_mask[subsets[-1]] = True
          torch.index_select(node_mask, 0, row, out=edge_mask)
          subsets.append(col[edge_mask])
  
      subset, inv = torch.cat(subsets).unique(return_inverse=True)
      inv = inv[: node_idx.numel()]
  
      node_mask.fill_(False)
      node_mask[subset] = True
  
      if not directed:
          edge_mask = node_mask[row] & node_mask[col]
  
      edge_index = edge_index[:, edge_mask]
  
      if relabel_nodes:
          mapping = row.new_full((num_nodes,), -1)
          mapping[subset] = torch.arange(
              subset.size(0), device=row.device
          )
          edge_index = mapping[edge_index]
  
      return subset, edge_index, inv, edge_mask
  
file_path: pytorch_geometric/torch_geometric/utils/_subgraph.py
imports: from typing import List, Literal, Optional, Tuple, Union, overload
  import torch
  from torch import Tensor
  from torch_geometric.typing import OptTensor, PairTensor
  from torch_geometric.utils import scatter
  from torch_geometric.utils.map import map_index
  from torch_geometric.utils.mask import index_to_mask
  from torch_geometric.utils.num_nodes import maybe_num_nodes
label_desc: Use image node_id 7 for calling a global function with example usage: k_hop_subgraph(node_idx, num_hops, edge_index, relabel_nodes, num_nodes, flow, directed) and returns: subset,  edge_index,  inv,  edge_mask