code: def _find_best_tree(edge_logits):
      K = len(edge_logits)
      V = int(round(0.5 + (0.25 + 2 * K) ** 0.5))
      assert K == V * (V - 1) // 2
      E = V - 1
      grid = make_complete_graph(V)
  
      # Each of E edges in the tree is stored as an id k in [0, K) indexing into
      # the complete graph. The id of an edge (v1,v2) is k = v1+v2*(v2-1)/2.
      edge_ids = torch.empty((E,), dtype=torch.long)
      # This maps each vertex to whether it is a member of the cumulative tree.
      components = torch.zeros(V, dtype=torch.bool)
  
      # Find the first edge.
      k = edge_logits.argmax(0).item()
      components[grid[:, k]] = 1
      edge_ids[0] = k
  
      # Find edges connecting the cumulative tree to a new leaf.
      for e in range(1, E):
          c1, c2 = components[grid]
          mask = c1 != c2
          valid_logits = edge_logits[mask]
          k = valid_logits.argmax(0).item()
          k = mask.nonzero(as_tuple=False)[k]
          components[grid[:, k]] = 1
          edge_ids[e] = k
  
      # Convert edge ids to a canonical list of pairs.
      edge_ids = edge_ids.sort()[0]
      edges = torch.empty((E, 2), dtype=torch.long)
      edges[:, 0] = grid[0, edge_ids]
      edges[:, 1] = grid[1, edge_ids]
      return edges
  
imports: import itertools
  import math
  import torch
  from torch.distributions import constraints
  from torch.distributions.utils import lazy_property
  from pyro.distributions.torch_distribution import TorchDistribution
label_desc: Use image node_id 12 for calling a global function with example usage: _find_best_tree(edge_logits) and returns: edges