text: import copy
  from typing import List, Optional, Tuple, NamedTuple
  
  import torch
  from torch_sparse import SparseTensor
  
  
  class Adj(NamedTuple):
      edge_index: torch.Tensor
      e_id: torch.Tensor
      size: Tuple[int, int]
  
      def to(self, *args, **kwargs):
          return Adj(self.edge_index.to(*args, **kwargs),
                     self.e_id.to(*args, **kwargs), self.size)
  
  
  class NeighborSampler(torch.utils.data.DataLoader):
      r"""The neighbor sampler from the `"Inductive Representation Learning on
      Large Graphs" <https://arxiv.org/abs/1706.02216>`_ paper, which allows
      for mini-batch training of GNNs on large-scale graphs where full-batch
      training is not feasible.
  
      Given a GNN with :math:`L` layers and a specific mini-batch of nodes
      :obj:`node_idx` for which we want to compute embeddings, this module
      iteratively samples neighbors and constructs bipartite graphs that simulate
      the actual computation flow of GNNs.
  
      More specifically, :obj:`sizes` denotes how much neighbors we want to
      sample for each node in each layer.
      This module then takes in these :obj:`sizes` and iteratively samples
      :obj:`sizes[l]` for each node involved in layer :obj:`l`.
      In the next layer, sampling is repeated for the union of nodes that were
      already encountered.
      The actual computation graphs are then returned in reverse-mode, meaning
      that we pass messages from a larger set of nodes to a smaller one, until we
      reach the nodes for which we originally wanted to compute embeddings.
  
      Hence, an item returned by :class:`NeighborSampler` holds the current
      :obj:`batch_size`, the IDs :obj:`n_id` of all nodes involved in the
      computation, and a list of bipartite graph objects via the tuple
      :obj:`(edge_index, e_id, size)`, where :obj:`edge_index` represents the
      bipartite edges between source and target nodes, :obj:`e_id` denotes the
      IDs of original edges in the full graph, and :obj:`size` holds the shape
      of the bipartite graph.
      For each bipartite graph, target nodes are also included at the beginning
      of the list of source nodes so that one can easily apply skip-connections
      or add self-loops.
  
      .. note::
  
          For an example of using :obj:`NeighborSampler`, see
          `examples/reddit.py
          <https://github.com/rusty1s/pytorch_geometric/blob/master/examples/
          reddit.py>`_ or
          `examples/ogbn_products_sage.py
          <https://github.com/rusty1s/pytorch_geometric/blob/master/examples/
          ogbn_products_sage.py>`_.
  
      Args:
          edge_index (LongTensor): The edge indices of the full-graph.
          size ([int]): The number of neighbors to
              sample for each node in each layer. If set to :obj:`sizes[i] = -1`,
              all neighbors are included in layer :obj:`l`.
          node_idx (LongTensor, optional): The nodes that should be considered
              for creating mini-batches. If set to :obj:`None`, all nodes will be
              considered.
          flow (string, optional): The flow direction of message passing
              (:obj:`"source_to_target"` or :obj:`"target_to_source"`).
              (default: :obj:`"source_to_target"`)
          **kwargs (optional): Additional arguments of
              :class:`torch.utils.data.DataLoader`, such as :obj:`batch_size`,
              :obj:`shuffle`, :obj:`drop_last` or :obj:`num_workers`.
      """
      def __init__(self, edge_index: torch.Tensor, sizes: List[int],
                   node_idx: Optional[torch.Tensor] = None,
                   num_nodes: Optional[int] = None,
                   flow: str = "source_to_target", **kwargs):
  
          edge_index = edge_index.cpu()
          N = int(edge_index.max() + 1) if num_nodes is None else num_nodes
          edge_attr = torch.arange(edge_index.size(1))
          adj = SparseTensor(row=edge_index[0], col=edge_index[1],
                             value=edge_attr, sparse_sizes=(N, N),
                             is_sorted=False)
          adj = adj.t() if flow == 'source_to_target' else adj
          self.adj = adj.to('cpu')
  
          if node_idx is None:
              node_idx = torch.arange(N)
          elif node_idx.dtype == torch.bool:
              node_idx = node_idx.nonzero(as_tuple=False).view(-1)
  
          self.sizes = sizes
          self.flow = flow
          assert self.flow in ['source_to_target', 'target_to_source']
  
          super(NeighborSampler, self).__init__(node_idx.tolist(),
                                                collate_fn=self.sample, **kwargs)
  
      def sample(self, batch):
          if not isinstance(batch, torch.Tensor):
              batch = torch.tensor(batch)
  
          batch_size: int = len(batch)
          adjs: List[Adj] = []
  
          n_id = batch
          for size in self.sizes:
              adj, n_id = self.adj.sample_adj(n_id, size, replace=False)
              if self.flow == 'source_to_target':
                  adj = adj.t()
              row, col, e_id = adj.coo()
              size = adj.sparse_sizes()
              edge_index = torch.stack([row, col], dim=0)
  
              adjs.append(Adj(edge_index, e_id, size))
  
          if len(adjs) > 1:
              return batch_size, n_id, adjs[::-1]
          else:
              return batch_size, n_id, adjs[0]
  
      def __repr__(self):
          return '{}(sizes={})'.format(self.__class__.__name__, self.sizes)
  
  
  class RandomIndexSampler(torch.utils.data.Sampler):
      def __init__(self, num_nodes: int, num_parts: int, shuffle: bool = False):
          self.N = num_nodes
          self.num_parts = num_parts
          self.shuffle = shuffle
          self.n_ids = self.get_node_indices()
  
      def get_node_indices(self):
          n_id = torch.randint(self.num_parts, (self.N, ), dtype=torch.long)
          n_ids = [(n_id == i).nonzero(as_tuple=False).view(-1)
                   for i in range(self.num_parts)]
          return n_ids
  
      def __iter__(self):
          if self.shuffle:
              self.n_ids = self.get_node_indices()
          return iter(self.n_ids)
  
      def __len__(self):
          return self.num_parts
  
  
  class RandomNodeSampler(torch.utils.data.DataLoader):
      r"""A data loader that randomly samples nodes within a graph and returns
      their induced subgraph.
  
      .. note::
  
          For an example of using :obj:`RandomNodeSampler`, see
          `examples/ogbn_proteins_deepgcn.py
          <https://github.com/rusty1s/pytorch_geometric/blob/master/examples/
          ogbn_proteins_deepgcn.py>`_.
  
      Args:
          data (torch_geometric.data.Data): The graph data object.
          num_parts (int): The number of partitions.
          shuffle (bool, optional): If set to :obj:`True`, the data is reshuffled
              at every epoch (default: :obj:`False`).
          **kwargs (optional): Additional arguments of
              :class:`torch.utils.data.DataLoader`, such as :obj:`num_workers`.
      """
      def __init__(self, data, num_parts: int, shuffle: bool = False, **kwargs):
          assert data.edge_index is not None
  
          self.N = N = data.num_nodes
          self.E = data.num_edges
  
          self.adj = SparseTensor(
              row=data.edge_index[0], col=data.edge_index[1],
              value=torch.arange(self.E, device=data.edge_index.device),
              sparse_sizes=(N, N))
  
          self.data = copy.copy(data)
          self.data.edge_index = None
  
          super(RandomNodeSampler, self).__init__(
              self, batch_size=1,
              sampler=RandomIndexSampler(self.N, num_parts, shuffle),
              collate_fn=self.__collate__, **kwargs)
  
      def __getitem__(self, idx):
          return idx
  
      def __collate__(self, node_idx):
          node_idx = node_idx[0]
  
          data = self.data.__class__()
          data.num_nodes = node_idx.size(0)
          adj, _ = self.adj.saint_subgraph(node_idx)
          row, col, edge_idx = adj.coo()
          data.edge_index = torch.stack([row, col], dim=0)
  
          for key, item in self.data:
              if isinstance(item, torch.Tensor) and item.size(0) == self.N:
                  data[key] = item[node_idx]
              elif isinstance(item, torch.Tensor) and item.size(0) == self.E:
                  data[key] = item[edge_idx]
              else:
                  data[key] = item
  
          return data
  