code: def monte_carlo_variational_loss(
      target_log_prob_fn,
      surrogate_posterior,
      sample_size=1,
      importance_sample_size=1,
      discrepancy_fn=kl_reverse,
      use_reparameterization=None,
      gradient_estimator=None,
      stopped_surrogate_posterior=None,
      seed=None,
      name=None,
  ):
      """Monte-Carlo approximation of an f-Divergence variational loss.
  
      Variational losses measure the divergence between an unnormalized target
      distribution `p` (provided via `target_log_prob_fn`) and a surrogate
      distribution `q` (provided as `surrogate_posterior`). When the
      target distribution is an unnormalized posterior from conditioning a model on
      data, minimizing the loss with respect to the parameters of
      `surrogate_posterior` performs approximate posterior inference.
  
      This function defines losses of the form
      `E_q[discrepancy_fn(log(u))]`, where `u = p(z) / q(z)` in the (default) case
      where `importance_sample_size == 1`, and
      `u = mean([p(z[k]) / q(z[k]) for k in range(importance_sample_size)]))` more
      generally. These losses are sometimes known as f-divergences [1, 2].
  
      The default behavior (`discrepancy_fn == tfp.vi.kl_reverse`, where
      `tfp.vi.kl_reverse = lambda logu: -logu`, and
      `importance_sample_size == 1`) computes an unbiased estimate of the standard
      evidence lower bound (ELBO) [3]. The bound may be tightened by setting
      `importance_sample_size > 1` [4], and the variance of the estimate reduced by
      setting `sample_size > 1`. Other discrepancies of interest
      available under `tfp.vi` include the forward `KL[p||q]`, total variation
      distance, Amari alpha-divergences, and [more](
      https://en.wikipedia.org/wiki/F-divergence).
  
      Args:
        target_log_prob_fn: Python callable that takes a set of `Tensor` arguments
          and returns a `Tensor` log-density. Given
          `q_sample = surrogate_posterior.sample(sample_size)`, this
          will be called as `target_log_prob_fn(*q_sample)` if `q_sample` is a list
          or a tuple, `target_log_prob_fn(**q_sample)` if `q_sample` is a
          dictionary, or `target_log_prob_fn(q_sample)` if `q_sample` is a `Tensor`.
          It should support batched evaluation, i.e., should return a result of
          shape `[sample_size]`.
        surrogate_posterior: A `tfp.distributions.Distribution`
          instance defining a variational posterior (could be a
          `tfd.JointDistribution`). If using `tf.Variable` parameters, the
          distribution's `log_prob` and (if reparameterizeable) `sample` methods
          must directly invoke all ops that generate gradients to the underlying
          variables. One way to ensure this is to use `tfp.util.TransformedVariable`
          and/or `tfp.util.DeferredTensor` to represent any parameters defined as
          transformations of unconstrained variables, so that the transformations
          execute at runtime instead of at distribution creation.
        sample_size: Integer scalar number of Monte Carlo samples used to
          approximate the variational divergence. Larger values may stabilize
          the optimization, but at higher cost per step in time and memory.
          Default value: `1`.
        importance_sample_size: Python `int` number of terms used to define an
          importance-weighted divergence. If `importance_sample_size > 1`, then the
          `surrogate_posterior` is optimized to function as an importance-sampling
          proposal distribution. In this case it often makes sense to use
          importance sampling to approximate posterior expectations (see
          `tfp.vi.fit_surrogate_posterior` for an example).
          Default value: `1`.
        discrepancy_fn: Python `callable` representing a Csiszar `f` function in
          in log-space. That is, `discrepancy_fn(log(u)) = f(u)`, where `f` is
          convex in `u`.
          Default value: `tfp.vi.kl_reverse`.
        use_reparameterization: Deprecated; use `gradient_estimator` instead.
        gradient_estimator: Optional element from `tfp.vi.GradientEstimators`
          specifying the stochastic gradient estimator to associate with the
          variational loss. If `None`, a default estimator (either score-function or
          reparameterization) is chosen based on
          `surrogate_posterior.reparameterization_type`.
          Default value: `None`.
        stopped_surrogate_posterior: Optional copy of `surrogate_posterior` with
          stopped gradients to the parameters, e.g.,
          `tfd.Normal(loc=tf.stop_gradient(loc), scale=tf.stop_gradient(scale))`.
          Required if and only if
          `gradient_estimator == tfp.vi.GradientEstimators.DOUBLY_REPARAMETERIZED`.
          Default value: `None`.
        seed: PRNG seed for `surrogate_posterior.sample`; see
          `tfp.random.sanitize_seed` for details.
        name: Python `str` name prefixed to Ops created by this function.
  
      Returns:
        monte_carlo_variational_loss: `float`-like `Tensor` Monte Carlo
          approximation of the Csiszar f-Divergence.
  
      Raises:
        ValueError: if `surrogate_posterior` is not a reparameterized
          distribution and `use_reparameterization = True`. A distribution is said
          to be "reparameterized" when its samples are generated by transforming the
          samples of another distribution that does not depend on the first
          distribution's parameters. This property ensures the gradient with respect
          to parameters is valid.
        TypeError: if `target_log_prob_fn` is not a Python `callable`.
  
      #### Csiszar f-divergences
  
      A Csiszar function `f` is a convex function from `R^+` (the positive reals)
      to `R`. The Csiszar f-Divergence is given by:
  
      ```none
      D_f[p(X), q(X)] := E_{q(X)}[ f( p(X) / q(X) ) ]
                      ~= m**-1 sum_j^m f( p(x_j) / q(x_j) ),
                                 where x_j ~iid q(X)
      ```
  
      For example, `f = lambda u: -log(u)` recovers `KL[q||p]`, while `f =
      lambda u: u * log(u)` recovers the forward `KL[p||q]`. These and other
      functions are available in `tfp.vi`.
  
      #### Example Application:
  
      The Csiszar f-Divergence is a useful framework for variational inference.
      I.e., observe that,
  
      ```none
      f(p(x)) =  f( E_{q(Z | x)}[ p(x, Z) / q(Z | x) ] )
              <= E_{q(Z | x)}[ f( p(x, Z) / q(Z | x) ) ]
              := D_f[p(x, Z), q(Z | x)]
      ```
  
      The inequality follows from the fact that the "perspective" of `f`, i.e.,
      `(s, t) |-> t f(s / t))`, is convex in `(s, t)` when `s/t in domain(f)` and
      `t` is a real. Since the above framework includes the popular Evidence Lower
      BOund (ELBO) as a special case, i.e., `f(u) = -log(u)`, we call this framework
      "Evidence Divergence Bound Optimization" (EDBO).
  
      #### References:
  
      [1]: https://en.wikipedia.org/wiki/F-divergence
  
      [2]: Ali, Syed Mumtaz, and Samuel D. Silvey. "A general class of coefficients
           of divergence of one distribution from another." Journal of the Royal
           Statistical Society: Series B (Methodological) 28.1 (1966): 131-142.
  
      [3]: Christopher M. Bishop. Pattern Recognition and Machine Learning.
           Springer, 2006.
  
      [4]  Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Importance Weighted
           Autoencoders. In _International Conference on Learning
           Representations_, 2016. https://arxiv.org/abs/1509.00519
  
      """
      with tf.name_scope(name or "monte_carlo_variational_loss"):
          if not callable(target_log_prob_fn):
              raise TypeError(
                  "`target_log_prob_fn` must be a Python `callable`"
                  "function."
              )
  
          sample_seed, target_seed = samplers.split_seed(seed, 2)
          reparameterization_types = tf.nest.flatten(
              surrogate_posterior.reparameterization_type
          )
          if gradient_estimator is None:
              gradient_estimator = _choose_gradient_estimator(
                  use_reparameterization=use_reparameterization,
                  reparameterization_types=reparameterization_types,
              )
  
          if gradient_estimator == GradientEstimators.VIMCO:
              return csiszar_vimco(
                  f=discrepancy_fn,
                  p_log_prob=target_log_prob_fn,
                  q=surrogate_posterior,
                  num_draws=importance_sample_size,
                  num_batch_draws=sample_size,
                  seed=seed,
              )
          if gradient_estimator == GradientEstimators.SCORE_FUNCTION:
              if tf.get_static_value(importance_sample_size) != 1:
                  # TODO(b/213378570): Support score function gradients for
                  # importance-weighted bounds.
                  raise ValueError(
                      "Score-function gradients are not supported for "
                      "losses with `importance_sample_size != 1`."
                  )
              # Score fn objective requires explicit gradients of `log_prob`.
              q_samples = surrogate_posterior.sample(
                  [sample_size * importance_sample_size],
                  seed=sample_seed,
              )
              q_lp = None
          else:
              if any(
                  reparameterization_type != FULLY_REPARAMETERIZED
                  for reparameterization_type in reparameterization_types
              ):
                  warnings.warn(
                      "Reparameterization gradients requested, but "
                      "`surrogate_posterior.reparameterization_type` is not fully "
                      "reparameterized (saw: {}). Gradient estimates may be "
                      "biased.".format(
                          surrogate_posterior.reparameterization_type
                      )
                  )
              # Attempt to avoid bijector inverses by computing the surrogate log prob
              # during the forward sampling pass.
              (
                  q_samples,
                  q_lp,
              ) = surrogate_posterior.experimental_sample_and_log_prob(
                  [sample_size * importance_sample_size],
                  seed=sample_seed,
              )
  
          return monte_carlo.expectation(
              f=_make_importance_weighted_divergence_fn(
                  target_log_prob_fn,
                  surrogate_posterior=surrogate_posterior,
                  discrepancy_fn=discrepancy_fn,
                  precomputed_surrogate_log_prob=q_lp,
                  importance_sample_size=importance_sample_size,
                  gradient_estimator=gradient_estimator,
                  stopped_surrogate_posterior=stopped_surrogate_posterior,
                  seed=target_seed,
              ),
              samples=q_samples,
              # Log-prob is only used if `gradient_estimator == SCORE_FUNCTION`.
              log_prob=surrogate_posterior.log_prob,
              use_reparameterization=(
                  gradient_estimator
                  != GradientEstimators.SCORE_FUNCTION
              ),
          )
  
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 19 for calling a global function with example usage: monte_carlo_variational_loss(target_log_prob_fn, surrogate_posterior, sample_size, importance_sample_size, discrepancy_fn, use_reparameterization, gradient_estimator, stopped_surrogate_posterior, seed, name) and returns: monte_carlo, csiszar_vimco