code: def _donsker_varadhan_loss(
      model, T, observation_labels, target_labels
  ):
      """DV loss: to evaluate directly use `donsker_varadhan_eig` setting `num_steps=0`."""
  
      ewma_log = EwmaLog(alpha=0.90)
  
      def loss_fn(design, num_particles, **kwargs):
          try:
              pyro.module("T", T)
          except AssertionError:
              pass
  
          expanded_design = lexpand(design, num_particles)
  
          # Unshuffled data
          unshuffled_trace = poutine.trace(model).get_trace(
              expanded_design
          )
          y_dict = {
              l: unshuffled_trace.nodes[l]["value"]
              for l in observation_labels
          }
  
          # Shuffled data
          # Not actually shuffling, resimulate for safety
          conditional_model = pyro.condition(model, data=y_dict)
          shuffled_trace = poutine.trace(conditional_model).get_trace(
              expanded_design
          )
  
          T_joint = T(
              expanded_design,
              unshuffled_trace,
              observation_labels,
              target_labels,
          )
          T_independent = T(
              expanded_design,
              shuffled_trace,
              observation_labels,
              target_labels,
          )
  
          joint_expectation = T_joint.sum(0) / num_particles
  
          A = T_independent - math.log(num_particles)
          s, _ = torch.max(A, dim=0)
          independent_expectation = s + ewma_log(
              (A - s).exp().sum(dim=0), s
          )
  
          loss = joint_expectation - independent_expectation
          # Switch sign, sum over batch dimensions for scalar loss
          agg_loss = -loss.sum()
          return agg_loss, loss
  
      return loss_fn
  
imports: import math
  import warnings
  import torch
  import pyro
  from pyro import poutine
  from pyro.contrib.oed.search import Search
  from pyro.contrib.util import lexpand
  from pyro.infer import SVI, EmpiricalMarginal, Importance
  from pyro.infer.autoguide.utils import mean_field_entropy
  from pyro.util import torch_isinf, torch_isnan
label_desc: Use image node_id 16 for calling a global function with example usage: _donsker_varadhan_loss(model, T, observation_labels, target_labels) and returns: loss_fn, agg_loss,  loss