code: def compute_marginals_persistent_bp(
      exists_logits, assign_logits, bp_iters, bp_momentum=0.5
  ):
      """
      This implements approximate inference of pairwise marginals via
      loopy belief propagation, adapting the approach of [1], [2].
  
      See :class:`MarginalAssignmentPersistent` 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
      [2] Ryan Turner, Steven Bottone, Bhargav Avasarala (2014)
          A Complete Variational Tracker
          https://papers.nips.cc/paper/5572-a-complete-variational-tracker.pdf
      """
      # This implements forward-backward message passing among three sets of variables:
      #
      #   a[t,j] ~ Categorical(num_objects + 1), detection -> object assignment
      #   b[t,i] ~ Categorical(num_detections + 1), object -> detection assignment
      #     e[i] ~ Bernonulli, whether each object exists
      #
      # Only assign = a and exists = e are returned.
      assert 0 <= bp_momentum < 1, bp_momentum
      old, new = bp_momentum, 1 - bp_momentum
      num_frames, num_detections, num_objects = assign_logits.shape
      dtype = assign_logits.dtype
      device = assign_logits.device
      message_b_to_a = torch.zeros(
          num_frames,
          num_detections,
          num_objects,
          dtype=dtype,
          device=device,
      )
      message_a_to_b = torch.zeros(
          num_frames,
          num_detections,
          num_objects,
          dtype=dtype,
          device=device,
      )
      message_b_to_e = torch.zeros(
          num_frames, num_objects, dtype=dtype, device=device
      )
      message_e_to_b = torch.zeros(
          num_frames, num_objects, dtype=dtype, device=device
      )
  
      for i in range(bp_iters):
          odds_a = (assign_logits + message_b_to_a).exp()
          message_a_to_b = old * message_a_to_b + new * (
              assign_logits - (odds_a.sum(2, True) - odds_a).log1p()
          )
          message_b_to_e = (
              old * message_b_to_e
              + new * message_a_to_b.exp().sum(1).log1p()
          )
          message_e_to_b = old * message_e_to_b + new * (
              exists_logits + message_b_to_e.sum(0) - message_b_to_e
          )
          odds_b = message_a_to_b.exp()
          message_b_to_a = (
              old * message_b_to_a
              - new
              * (
                  (-message_e_to_b).exp().unsqueeze(1)
                  + (1 + odds_b.sum(1, True) - odds_b)
              ).log()
          )
  
          warn_if_nan(
              message_a_to_b, "message_a_to_b iter {}".format(i)
          )
          warn_if_nan(
              message_b_to_e, "message_b_to_e iter {}".format(i)
          )
          warn_if_nan(
              message_e_to_b, "message_e_to_b iter {}".format(i)
          )
          warn_if_nan(
              message_b_to_a, "message_b_to_a iter {}".format(i)
          )
  
      # Convert from probs to logits.
      exists = exists_logits + message_b_to_e.sum(0)
      assign = assign_logits + message_b_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 7 for calling a global function with example usage: compute_marginals_persistent_bp(exists_logits, assign_logits, bp_iters, bp_momentum) and returns: exists,  assign