class_docstr: Description
  -----------
  GraphSAINT with edge sampler.
  
  Parameters
  ----------
  edge_budget : int
      the expected number of edges in each subgraph, which is specifically explained in the paper.
class_docstr_tok: ["Description","--","--","--","--","--","-","GraphSAINT","with","edge","sampler",".","Parameters","--","--","--","--","--","edge_budget",":","int","the","expected","number","of","edges","in","each","subgraph",",","which","is","specifically","explained","in","the","paper","."]
code: def __sample__(self):
      if self.prob is None:
          src, dst = self.train_g.edges()
          src_degrees, dst_degrees = self.train_g.in_degrees(
              src
          ).float().clamp(min=1), self.train_g.in_degrees(
              dst
          ).float().clamp(
              min=1
          )
          prob_mat = 1.0 / src_degrees + 1.0 / dst_degrees
          prob_mat = scipy.sparse.csr_matrix(
              (prob_mat.numpy(), (src.numpy(), dst.numpy()))
          )
          # The edge probability here only contains that of edges in upper triangle adjacency matrix
          # Because we assume the graph is undirected, that is, the adjacency matrix is symmetric. We only need
          # to consider half of edges in the graph.
          self.prob = th.tensor(scipy.sparse.triu(prob_mat).data)
          self.prob /= self.prob.sum()
          self.adj_nodes = np.stack(prob_mat.nonzero(), axis=1)
  
      sampled_edges = np.unique(
          dgl.random.choice(
              len(self.prob),
              size=self.edge_budget,
              prob=self.prob,
              replace=False,
          )
      )
      sampled_nodes = np.unique(
          self.adj_nodes[sampled_edges].flatten()
      ).astype("long")
      return sampled_nodes
  
imports: import math
  import os
  import random
  import time
  import dgl
  import dgl.function
  import numpy
  import scipy
  import torch
  from dgl.sampling import pack_traces, random_walk
  from torch.utils.data import DataLoader
label_desc: Use image node_id 3 for calling the SAINTEdgeSampler obj's underlying member method code with example usage: obj.__sample__() and returns: sampled_nodes