code: def compute_marginals_sparse_bp(
      num_objects,
      num_detections,
      edges,
      exists_logits,
      assign_logits,
      bp_iters,
  ):
      """
      This implements approximate inference of pairwise marginals via
      loopy belief propagation, adapting the approach of [1].
  
      See :class:`MarginalAssignmentSparse` for args and problem description.
  
      [1] Jason L. Williams, Roslyn A. Lau (2014)
          Approximate evaluation of marginal association probabilities with
          belief propagation
          https://arxiv.org/abs/1209.6299
      """
      exists_factor = exists_logits[edges[1]]
  
      def sparse_sum(x, dim, keepdim=False):
          assert dim in (0, 1)
          x = torch.zeros(
              [num_objects, num_detections][dim],
              dtype=x.dtype,
              device=x.device,
          ).scatter_add_(0, edges[1 - dim], x)
          if keepdim:
              x = x[edges[1 - dim]]
          return x
  
      message_e_to_a = torch.zeros_like(assign_logits)
      message_a_to_e = torch.zeros_like(assign_logits)
      for i in range(bp_iters):
          message_e_to_a = (
              -(
                  message_a_to_e
                  - sparse_sum(message_a_to_e, 0, True)
                  - exists_factor
              )
              .exp()
              .log1p()
          )
          joint = (assign_logits + message_e_to_a).exp()
          message_a_to_e = (
              (
                  assign_logits
                  - torch.log1p(sparse_sum(joint, 1, True) - joint)
              )
              .exp()
              .log1p()
          )
          warn_if_nan(
              message_e_to_a, "message_e_to_a iter {}".format(i)
          )
          warn_if_nan(
              message_a_to_e, "message_a_to_e iter {}".format(i)
          )
  
      # Convert from probs to logits.
      exists = exists_logits + sparse_sum(message_a_to_e, 0)
      assign = assign_logits + message_e_to_a
      warn_if_nan(exists, "exists")
      warn_if_nan(assign, "assign")
      return exists, assign
  
file_path: pyro/pyro/contrib/tracking/assignment.py
imports: import itertools
  import math
  import numbers
  import torch
  import pyro.distributions
  from pyro.util import warn_if_nan
label_desc: Use image node_id 5 for calling a global function with example usage: compute_marginals_sparse_bp(num_objects, num_detections, edges, exists_logits, assign_logits, bp_iters) and returns: exists,  assign, x