class_docstr: Common methods for testing ADVI surrogate posteriors.
class_docstr_tok: ["Common","methods","for","testing","ADVI","surrogate","posteriors","."]
docstr: Initializes a stateless (JAX) or stateful (TF) surrogate posterior.
code: def _initialize_surrogate(
      self, family_str, seed, is_stateless=JAX_MODE, **kwargs
  ):
      """Initializes a stateless (JAX) or stateful (TF) surrogate posterior."""
      # TODO(davmre): refactor to support testing gradients of stateless
      # surrogates.
      if is_stateless:
          build_fn = getattr(
              surrogate_posteriors, family_str + "_stateless"
          )
          init_fn, apply_fn = build_fn(**kwargs)
          return apply_fn(init_fn(seed=seed))
      else:
          build_fn = getattr(surrogate_posteriors, family_str)
          surrogate_posterior = build_fn(seed=seed, **kwargs)
          self.evaluate(
              [
                  v.initializer
                  for v in surrogate_posterior.trainable_variables
              ]
          )
          return surrogate_posterior
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import exp
  from tensorflow_probability.python.bijectors import fill_triangular
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.bijectors import joint_map
  from tensorflow_probability.python.bijectors import masked_autoregressive
  from tensorflow_probability.python.bijectors import sigmoid
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.distributions import deterministic
  from tensorflow_probability.python.distributions import exponential
  from tensorflow_probability.python.distributions import gamma
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import laplace
  from tensorflow_probability.python.distributions import logistic
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import student_t
  from tensorflow_probability.python.experimental.vi import surrogate_posteriors
  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 import test_util
  from tensorflow_probability.python.internal import tf_keras
  from tensorflow_probability.python.vi import optimization
  from tensorflow.python.util import nest
label_desc: Use image node_id 1 for calling the _SurrogatePosterior obj's underlying member method code with example usage: obj._initialize_surrogate(family_str, seed, is_stateless) and returns: apply_fn, surrogate_posterior