code: def _factored_surrogate_posterior(  # pylint: disable=dangerous-default-value
      event_shape=None,
      bijector=None,
      batch_shape=(),
      base_distribution_cls=normal.Normal,
      initial_parameters={"scale": 1e-2},
      dtype=tf.float32,
      validate_args=False,
      name=None,
  ):
      """Builds a joint variational posterior that factors over model variables.
  
      By default, this method creates an independent trainable Normal distribution
      for each variable, transformed using a bijector (if provided) to
      match the support of that variable. This makes extremely strong
      assumptions about the posterior: that it is approximately normal (or
      transformed normal), and that all model variables are independent.
  
      Args:
        event_shape: `Tensor` shape, or nested structure of `Tensor` shapes,
          specifying the event shape(s) of the posterior variables.
        bijector: Optional `tfb.Bijector` instance, or nested structure of such
          instances, defining support(s) of the posterior variables. The structure
          must match that of `event_shape` and may contain `None` values. A
          posterior variable will be modeled as
          `tfd.TransformedDistribution(underlying_dist, bijector)` if a
          corresponding constraining bijector is specified, otherwise it is modeled
          as supported on the unconstrained real line.
        batch_shape: The `batch_shape` of the output distribution.
          Default value: `()`.
        base_distribution_cls: Subclass of `tfd.Distribution` that is instantiated
          and optionally transformed by the bijector to define the component
          distributions. May optionally be a structure of such subclasses
          matching `event_shape`.
          Default value: `tfd.Normal`.
        initial_parameters: Optional `str : Tensor` dictionary specifying initial
          values for some or all of the base distribution's trainable parameters,
          or a Python `callable` with signature
          `value = parameter_init_fn(parameter_name, shape, dtype, seed,
          constraining_bijector)`, passed to `tfp.experimental.util.make_trainable`.
          May optionally be a structure matching `event_shape` of such dictionaries
          and/or callables. Dictionary entries that do not correspond to parameter
          names are ignored.
          Default value: `{'scale': 1e-2}` (ignored when `base_distribution` does
            not have a `scale` parameter).
        dtype: Optional float `dtype` for trainable parameters. May
          optionally be a structure of such `dtype`s matching `event_shape`.
          Default value: `tf.float32`.
        validate_args: Python `bool`. Whether to validate input with asserts. This
          imposes a runtime cost. If `validate_args` is `False`, and the inputs are
          invalid, correct behavior is not guaranteed.
          Default value: `False`.
        name: Python `str` name prefixed to ops created by this function.
          Default value: `None` (i.e., 'build_factored_surrogate_posterior').
      Yields:
        *parameters: sequence of `trainable_state_util.Parameter` namedtuples.
          These are intended to be consumed by
          `trainable_state_util.as_stateful_builder` and
          `trainable_state_util.as_stateless_builder` to define stateful and
          stateless variants respectively.
  
      ### Examples
  
      Consider a Gamma model with unknown parameters, expressed as a joint
      Distribution:
  
      ```python
      Root = tfd.JointDistributionCoroutine.Root
      def model_fn():
        concentration = yield Root(tfd.Exponential(1.))
        rate = yield Root(tfd.Exponential(1.))
        y = yield tfd.Sample(tfd.Gamma(concentration=concentration, rate=rate),
                             sample_shape=4)
      model = tfd.JointDistributionCoroutine(model_fn)
      ```
  
      Let's use variational inference to approximate the posterior over the
      data-generating parameters for some observed `y`. We'll build a
      surrogate posterior distribution by specifying the shapes of the latent
      `rate` and `concentration` parameters, and that both are constrained to
      be positive.
  
      ```python
      surrogate_posterior = tfp.experimental.vi.build_factored_surrogate_posterior(
        event_shape=model.event_shape_tensor()[:-1],  # Omit the observed `y`.
        bijector=[tfb.Softplus(),   # Rate is positive.
                  tfb.Softplus()])  # Concentration is positive.
      ```
  
      This creates a trainable joint distribution, defined by variables in
      `surrogate_posterior.trainable_variables`. We use `fit_surrogate_posterior`
      to fit this distribution by minimizing a divergence to the true posterior.
  
      ```python
      y = [0.2, 0.5, 0.3, 0.7]
      losses = tfp.vi.fit_surrogate_posterior(
        lambda rate, concentration: model.log_prob([rate, concentration, y]),
        surrogate_posterior=surrogate_posterior,
        num_steps=100,
        optimizer=tf_keras.optimizers.Adam(0.1),
        sample_size=10)
  
      # After optimization, samples from the surrogate will approximate
      # samples from the true posterior.
      samples = surrogate_posterior.sample(100)
      posterior_mean = [tf.reduce_mean(x) for x in samples]     # mean ~= [1.1, 2.1]
      posterior_std = [tf.math.reduce_std(x) for x in samples]  # std  ~= [0.3, 0.8]
      ```
  
      If we wanted to initialize the optimization at a specific location, we can
      specify initial parameters when we build the surrogate posterior. Note that
      these parameterize the distribution(s) over unconstrained values,
      so we need to transform our desired constrained locations using the inverse
      of the constraining bijector(s).
  
      ```python
      surrogate_posterior = tfp.experimental.vi.build_factored_surrogate_posterior(
        event_shape=tf.nest.map_fn(tf.shape, initial_loc),
        bijector={'concentration': tfb.Softplus(),   # Rate is positive.
                  'rate': tfb.Softplus()}   # Concentration is positive.
        initial_parameters={
          'concentration': {'loc': tfb.Softplus().inverse(0.4), 'scale': 1e-2},
          'rate': {'loc': tfb.Softplus().inverse(0.2), 'scale': 1e-2}})
      ```
  
      """
      with tf.name_scope(name or "build_factored_surrogate_posterior"):
          # Convert event shapes to Tensors.
          shallow_structure = _get_event_shape_shallow_structure(
              event_shape
          )
          event_shape = nest.map_structure_up_to(
              shallow_structure,
              lambda s: ps.convert_to_shape_tensor(s, dtype=tf.int32),
              event_shape,
          )
  
          if nest.is_nested(bijector):
              event_space_bijector = joint_map.JointMap(
                  nest.map_structure(
                      lambda b: identity.Identity() if b is None else b,
                      nest_util.coerce_structure(event_shape, bijector),
                  ),
                  validate_args=validate_args,
              )
          else:
              event_space_bijector = bijector
  
          if event_space_bijector is None:
              unconstrained_event_shape = event_shape
          else:
              unconstrained_event_shape = (
                  event_space_bijector.inverse_event_shape_tensor(
                      event_shape
                  )
              )
          unconstrained_batch_and_event_shape = tf.nest.map_structure(
              lambda s: ps.concat([batch_shape, s], axis=0),
              unconstrained_event_shape,
          )
  
          base_distribution_cls = nest_util.broadcast_structure(
              event_shape, base_distribution_cls
          )
          try:
              # Check that we have initial parameters for each event part.
              nest.assert_shallow_structure(
                  event_shape, initial_parameters
              )
          except (ValueError, TypeError):
              # If not, broadcast the parameters to match the event structure.
              # We do this manually rather than using `nest_util.broadcast_structure`
              # because the initial parameters can themselves be structures (dicts).
              initial_parameters = nest.map_structure(
                  lambda x: initial_parameters, event_shape
              )
  
          unconstrained_trainable_distributions = yield from (
              nest_util.map_structure_coroutine(
                  trainable._make_trainable,  # pylint: disable=protected-access
                  cls=base_distribution_cls,
                  initial_parameters=initial_parameters,
                  batch_and_event_shape=unconstrained_batch_and_event_shape,
                  parameter_dtype=nest_util.broadcast_structure(
                      event_shape, dtype
                  ),
                  _up_to=event_shape,
              )
          )
          unconstrained_trainable_distribution = joint_distribution_util.independent_joint_distribution_from_structure(
              unconstrained_trainable_distributions,
              batch_ndims=ps.rank_from_shape(batch_shape),
              validate_args=validate_args,
          )
          if event_space_bijector is None:
              return unconstrained_trainable_distribution
          return transformed_distribution.TransformedDistribution(
              unconstrained_trainable_distribution, event_space_bijector
          )
  
imports: import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import chain
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.bijectors import invert
  from tensorflow_probability.python.bijectors import joint_map
  from tensorflow_probability.python.bijectors import reshape
  from tensorflow_probability.python.bijectors import restructure
  from tensorflow_probability.python.bijectors import scale_matvec_linear_operator
  from tensorflow_probability.python.bijectors import shift
  from tensorflow_probability.python.bijectors import split
  from tensorflow_probability.python.distributions import batch_broadcast
  from tensorflow_probability.python.distributions import joint_distribution_util
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.experimental.util import trainable
  from tensorflow_probability.python.experimental.vi.util import trainable_linear_operators
  from tensorflow_probability.python.internal import distribution_util
  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 import trainable_state_util
  from tensorflow.python.util import nest
label_desc: Use image node_id 2 for calling a global function with example usage: _factored_surrogate_posterior(event_shape, bijector, batch_shape, base_distribution_cls, initial_parameters, dtype, validate_args, name) and returns: transformed_distribution, unconstrained_trainable_distribution