code: def _make_importance_weighted_divergence_fn(
      target_log_prob_fn,
      surrogate_posterior,
      discrepancy_fn,
      precomputed_surrogate_log_prob=None,
      importance_sample_size=1,
      gradient_estimator=GradientEstimators.REPARAMETERIZATION,
      stopped_surrogate_posterior=None,
      seed=None,
  ):
      """Defines a function to compute an importance-weighted divergence."""
  
      def divergence_fn(q_samples):
          q_lp = precomputed_surrogate_log_prob
          target_log_prob = _call_fn_maybe_with_seed(
              target_log_prob_fn, q_samples, seed=seed
          )
  
          if (
              gradient_estimator
              == GradientEstimators.DOUBLY_REPARAMETERIZED
          ):
              # Sticking-the-landing is the special case of doubly-reparameterized
              # gradients with `importance_sample_size=1`.
              q_lp = stopped_surrogate_posterior.log_prob(q_samples)
          else:
              if q_lp is None:
                  q_lp = surrogate_posterior.log_prob(q_samples)
          log_weights = target_log_prob - q_lp
          return discrepancy_fn(log_weights)
  
      def importance_weighted_divergence_fn(q_samples):
          q_lp = precomputed_surrogate_log_prob
          if q_lp is None:
              q_lp = surrogate_posterior.log_prob(q_samples)
          target_log_prob = _call_fn_maybe_with_seed(
              target_log_prob_fn, q_samples, seed=seed
          )
          log_weights = target_log_prob - q_lp
  
          # Explicitly break out `importance_sample_size` as a separate axis.
          log_weights = tf.reshape(
              log_weights,
              ps.concat(
                  [
                      [-1, importance_sample_size],
                      ps.shape(log_weights)[1:],
                  ],
                  axis=0,
              ),
          )
          log_sum_weights = tf.reduce_logsumexp(log_weights, axis=1)
          log_avg_weights = log_sum_weights - tf.math.log(
              tf.cast(importance_sample_size, dtype=log_weights.dtype)
          )
  
          if (
              gradient_estimator
              == GradientEstimators.DOUBLY_REPARAMETERIZED
          ):
              # Adapted from original implementation at
              # https://github.com/google-research/google-research/blob/master/dreg_estimators/model.py
              normalized_weights = tf.stop_gradient(
                  tf.nn.softmax(log_weights, axis=1)
              )
              log_weights_with_stopped_q = tf.reshape(
                  target_log_prob
                  - stopped_surrogate_posterior.log_prob(q_samples),
                  ps.shape(log_weights),
              )
              dreg_objective = tf.reduce_sum(
                  log_weights_with_stopped_q
                  * tf.square(normalized_weights),
                  axis=1,
              )
              # Replace the objective's gradient with the doubly-reparameterized
              # gradient.
              log_avg_weights = tf.stop_gradient(log_avg_weights) + (
                  dreg_objective - tf.stop_gradient(dreg_objective)
              )
  
          return discrepancy_fn(log_avg_weights)
  
      if tf.get_static_value(importance_sample_size) == 1:
          return divergence_fn
      return importance_weighted_divergence_fn
  
imports: import enum
  import functools
  import warnings
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python import monte_carlo
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import nest_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal.reparameterization import FULLY_REPARAMETERIZED
  from tensorflow_probability.python.stats.leave_one_out import log_soomean_exp
  from tensorflow.python.util import deprecation
label_desc: Use image node_id 20 for calling a global function with example usage: _make_importance_weighted_divergence_fn(target_log_prob_fn, surrogate_posterior, discrepancy_fn, precomputed_surrogate_log_prob, importance_sample_size, gradient_estimator, stopped_surrogate_posterior, seed) and returns: importance_weighted_divergence_fn, discrepancy_fn, discrepancy_fn, divergence_fn