code: def test_edge_mean_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
      probs = d.log_prob(support).exp()[:, None].expand_as(k)
      expected = torch.zeros(K).scatter_add_(
          0, k.reshape(-1), probs.reshape(-1)
      )
  
      actual = d.edge_mean
      assert actual.shape == (V, V)
      v1, v2 = make_complete_graph(V)
      assert (actual[v1, v2] - expected).abs().max() < 1e-5, (
          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 8 for calling a global function with example usage: test_edge_mean_function(num_edges) without return types