code: def subgraph(
      subset: Union[Tensor, List[int]],
      edge_index: Tensor,
      edge_attr: OptTensor = None,
      relabel_nodes: bool = False,
      num_nodes: Optional[int] = None,
      *,
      return_edge_mask: bool = False,
  ) -> Union[
      Tuple[Tensor, OptTensor], Tuple[Tensor, OptTensor, Tensor]
  ]:
      r"""Returns the induced subgraph of :obj:`(edge_index, edge_attr)`
      containing the nodes in :obj:`subset`.
  
      Args:
          subset (LongTensor, BoolTensor or [int]): The nodes to keep.
          edge_index (LongTensor): The edge indices.
          edge_attr (Tensor, optional): Edge weights or multi-dimensional
              edge features. (default: :obj:`None`)
          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(edge_index) + 1`. (default: :obj:`None`)
          return_edge_mask (bool, optional): If set to :obj:`True`, will return
              the edge mask to filter out additional edge features.
              (default: :obj:`False`)
  
      :rtype: (:class:`LongTensor`, :class:`Tensor`)
  
      Examples:
          >>> edge_index = torch.tensor([[0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6],
          ...                            [1, 0, 2, 1, 3, 2, 4, 3, 5, 4, 6, 5]])
          >>> edge_attr = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12])
          >>> subset = torch.tensor([3, 4, 5])
          >>> subgraph(subset, edge_index, edge_attr)
          (tensor([[3, 4, 4, 5],
                  [4, 3, 5, 4]]),
          tensor([ 7.,  8.,  9., 10.]))
  
          >>> subgraph(subset, edge_index, edge_attr, return_edge_mask=True)
          (tensor([[3, 4, 4, 5],
                  [4, 3, 5, 4]]),
          tensor([ 7.,  8.,  9., 10.]),
          tensor([False, False, False, False, False, False,  True,
                  True,  True,  True,  False, False]))
      """
      device = edge_index.device
  
      if isinstance(subset, (list, tuple)):
          subset = torch.tensor(subset, dtype=torch.long, device=device)
  
      if subset.dtype != torch.bool:
          num_nodes = maybe_num_nodes(edge_index, num_nodes)
          node_mask = index_to_mask(subset, size=num_nodes)
      else:
          num_nodes = subset.size(0)
          node_mask = subset
          subset = node_mask.nonzero().view(-1)
  
      edge_mask = node_mask[edge_index[0]] & node_mask[edge_index[1]]
      edge_index = edge_index[:, edge_mask]
      edge_attr = (
          edge_attr[edge_mask] if edge_attr is not None else None
      )
  
      if relabel_nodes:
          edge_index, _ = map_index(
              edge_index.view(-1),
              subset,
              max_index=num_nodes,
              inclusive=True,
          )
          edge_index = edge_index.view(2, -1)
  
      if return_edge_mask:
          return edge_index, edge_attr, edge_mask
      else:
          return edge_index, edge_attr
  
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 5 for calling a global function with example usage: subgraph(subset, edge_index, edge_attr, relabel_nodes, num_nodes) and returns: edge_index,  edge_attr,  edge_mask, edge_index,  edge_attr