code: def log_mean_prob(trace, particle_dim):
      """
      Marginalizes out particle_dim from a trace.
      """
      assert particle_dim < 0
      trace.compute_log_prob()
      total = 0.0
      for node in trace.nodes.values():
          if (
              node["type"] == "sample"
              and type(node["fn"]).__name__ != "_Subsample"
          ):
              log_prob = node["log_prob"]
              assert log_prob.dim() == -particle_dim
              num_particles = log_prob.size(0)
              total = total + log_prob.reshape(num_particles, -1).sum(
                  -1
              )
      return total.logsumexp(0) - math.log(num_particles)
  
imports: import logging
  import math
  import warnings
  import pytest
  import torch
  from torch.autograd import grad
  import pyro
  import pyro.distributions
  from pyro import poutine
  from pyro.infer import TraceEnum_ELBO
  from pyro.infer.discrete import TraceEnumSample_ELBO, infer_discrete
  from pyro.infer.enum import config_enumerate
  from tests.common import assert_equal
label_desc: Use image node_id 2 for calling a global function with example usage: log_mean_prob(trace, particle_dim) and returns: unknown