text: import torch
  
  from .num_nodes import maybe_num_nodes
  
  
  def subgraph(subset, edge_index, edge_attr=None, relabel_nodes=False,
               num_nodes=None):
      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_val + 1` of :attr:`edge_index`. (default: :obj:`None`)
  
      :rtype: (:class:`LongTensor`, :class:`Tensor`)
      """
  
      device = edge_index.device
  
      if isinstance(subset, list) or isinstance(subset, tuple):
          subset = torch.tensor(subset, dtype=torch.long)
  
      if subset.dtype == torch.bool or subset.dtype == torch.uint8:
          n_mask = subset
  
          if relabel_nodes:
              n_idx = torch.zeros(
                  n_mask.size(0), dtype=torch.long, device=device)
              n_idx[subset] = torch.arange(subset.sum().item(), device=device)
      else:
          num_nodes = maybe_num_nodes(edge_index, num_nodes)
          n_mask = torch.zeros(num_nodes, dtype=torch.bool)
          n_mask[subset] = 1
  
          if relabel_nodes:
              n_idx = torch.zeros(num_nodes, dtype=torch.long, device=device)
              n_idx[subset] = torch.arange(subset.size(0), device=device)
  
      mask = n_mask[edge_index[0]] & n_mask[edge_index[1]]
      edge_index = edge_index[:, mask]
      edge_attr = edge_attr[mask] if edge_attr is not None else None
  
      if relabel_nodes:
          edge_index = n_idx[edge_index]
  
      return edge_index, edge_attr
  
  
  def k_hop_subgraph(node_idx, num_hops, edge_index, relabel_nodes=False,
                     num_nodes=None, flow='source_to_target'):
      r"""Computes the :math:`k`-hop subgraph of :obj:`edge_index` around node
      :attr:`node_idx`.
      It returns (1) the nodes involved in the subgraph, (2) the filtered
      :obj`edge_index` connectivity, and (3) the edge mask indicating which edges
      were preserved.
  
      Args:
          node_idx (int): The central node.
          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 (string, optional): The flow direction of :math:`k`-hop
              aggregation (:obj:`"source_to_target"` or
              :obj:`"target_to_source"`). (default: :obj:`"source_to_target"`)
  
      :rtype: (:class:`LongTensor`, :class:`LongTensor`, :class:`BoolTensor`)
      """
  
      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)
  
      subsets = [torch.tensor([node_idx], device=row.device).flatten()]
      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 = torch.cat(subsets).unique()
      # Add `node_idx` to the beginning of `subset`.
      subset = subset[subset != node_idx]
      subset = torch.cat([torch.tensor([node_idx], device=row.device), subset])
  
      node_mask.fill_(False)
      node_mask[subset] = True
      edge_mask = node_mask[row] & node_mask[col]
  
      edge_index = edge_index[:, edge_mask]
  
      if relabel_nodes:
          node_idx = row.new_full((num_nodes, ), -1)
          node_idx[subset] = torch.arange(subset.size(0), device=row.device)
          edge_index = node_idx[edge_index]
  
      return subset, edge_index, edge_mask
  