code: def sample_edge_neighborhood(
      adj_list, degrees, n_triplets, sample_size
  ):
      """Sample edges by neighborhool expansion.
  
      This guarantees that the sampled edges form a connected graph, which
      may help deeper GNNs that require information from more than one hop.
      """
      edges = np.zeros((sample_size), dtype=np.int32)
  
      # initialize
      sample_counts = np.array([d for d in degrees])
      picked = np.array([False for _ in range(n_triplets)])
      seen = np.array([False for _ in degrees])
  
      for i in range(0, sample_size):
          weights = sample_counts * seen
  
          if np.sum(weights) == 0:
              weights = np.ones_like(weights)
              weights[np.where(sample_counts == 0)] = 0
  
          probabilities = (weights) / np.sum(weights)
          chosen_vertex = np.random.choice(
              np.arange(degrees.shape[0]), p=probabilities
          )
          chosen_adj_list = adj_list[chosen_vertex]
          seen[chosen_vertex] = True
  
          chosen_edge = np.random.choice(
              np.arange(chosen_adj_list.shape[0])
          )
          chosen_edge = chosen_adj_list[chosen_edge]
          edge_number = chosen_edge[0]
  
          while picked[edge_number]:
              chosen_edge = np.random.choice(
                  np.arange(chosen_adj_list.shape[0])
              )
              chosen_edge = chosen_adj_list[chosen_edge]
              edge_number = chosen_edge[0]
  
          edges[i] = edge_number
          other_vertex = chosen_edge[1]
          picked[edge_number] = True
          sample_counts[chosen_vertex] -= 1
          sample_counts[other_vertex] -= 1
          seen[other_vertex] = True
  
      return edges
  
file_path: dgl/examples/tensorflow/rgcn/utils.py
imports: import dgl
  import numpy
  import tensorflow
label_desc: Use image node_id 2 for calling a global function with example usage: sample_edge_neighborhood(adj_list, degrees, n_triplets, sample_size) and returns: edges