code: def test_partition_function(num_edges):
      pyro.set_rng_seed(2**32 - num_edges)
      E = num_edges
      V = 1 + E
      K = V * (V - 1) // 2
      edge_logits = torch.randn(K)
      d = SpanningTree(edge_logits)
      with xfail_if_not_implemented():
          support = d.enumerate_support()
      v1 = support[..., 0]
      v2 = support[..., 1]
      k = v1 + v2 * (v2 - 1) // 2
      expected = edge_logits[k].sum(-1).logsumexp(0)
      actual = d.log_partition_function
      assert (actual - expected).abs() < 1e-6, (actual, expected)
  
file_path: pyro/tests/distributions/test_spanning_tree.py
imports: import logging
  import os
  from collections import Counter
  import pytest
  import torch
  import pyro
  from pyro.distributions.spanning_tree import NUM_SPANNING_TREES, SpanningTree, find_best_tree, make_complete_graph, sample_tree
  from tests.common import assert_equal, xfail_if_not_implemented
label_desc: Use image node_id 6 for calling a global function with example usage: test_partition_function(num_edges) without return types