code: def _asvi_convex_update_for_base_distribution(
      dist,
      mean_field=False,
      initial_prior_weight=0.5,
      sample_shape=None,
  ):
      """Creates a trainable surrogate for a (non-meta, non-joint) distribution."""
      posterior_batch_shape = dist.batch_shape_tensor()
      if sample_shape is not None:
          posterior_batch_shape = ps.concat(
              [
                  posterior_batch_shape,
                  distribution_util.expand_to_vector(sample_shape),
              ],
              axis=0,
          )
  
      temp_params_dict = {"name": _get_name(dist)}
      all_parameter_properties = dist.parameter_properties(
          dtype=dist.dtype
      )
      for param, prior_value in dist.parameters.items():
          if (
              param in (_NON_STATISTICAL_PARAMS + _NON_TRAINABLE_PARAMS)
              or prior_value is None
          ):
              temp_params_dict[param] = prior_value
              continue
  
          param_properties = all_parameter_properties[param]
          try:
              bijector = (
                  param_properties.default_constraining_bijector_fn()
              )
          except NotImplementedError:
              bijector = identity.Identity()
  
          param_shape = ps.concat(
              [
                  posterior_batch_shape,
                  ps.shape(prior_value)[
                      ps.rank(prior_value)
                      - param_properties.event_ndims :
                  ],
              ],
              axis=0,
          )
          # We need to pick a dtype that is compatible with the other parameters that
          # this parameter interacts with. We don't have direct access to this, so we
          # do some heuristics here.
          converted_prior_value = getattr(dist, param, None)
          if converted_prior_value is None:
              converted_prior_value = getattr(dist, "_" + param, None)
  
          if converted_prior_value is not None:
              # This is the safest option, if we can grab its converted value, then that
              # value is guaranteed to be good. Small chance that the implementor added
              # a property/attribute with the same name, but unrelated semantics...
              param_dtype = converted_prior_value.dtype
          elif bijector.dtype is not None:
              # For many bijectors, parameter_properties is implemented to specify a
              # good dtype value.
              param_dtype = bijector.dtype
          elif dtype_util.is_floating(dist.dtype):
              # If the distribution returns floats, all floats are probably of the same
              # type. This doesn't work for discrete distributions.
              param_dtype = dist.dtype
          else:
              # The greedy local conversion is unlikely to work, but we've run out of
              # places to look.
              param_dtype = tf.convert_to_tensor(prior_value).dtype
  
          # Initialize the mean-field parameter as a (constrained) standard
          # normal sample.
          # pylint: disable=cell-var-from-loop
          # Safe because the state utils guarantee to either call `init_fn`
          # immediately upon yielding, or not at all.
          mean_field_parameter = yield trainable_state_util.Parameter(
              init_fn=lambda seed: (  # pylint: disable=g-long-lambda
                  bijector.forward(
                      samplers.normal(
                          shape=bijector.inverse_event_shape(
                              param_shape
                          ),
                          dtype=bijector.inverse_dtype(param_dtype),
                          seed=seed,
                      )
                  )
              ),
              name="mean_field_parameter_{}_{}".format(
                  _get_name(dist), param
              ),
              constraining_bijector=bijector,
          )
          if mean_field:
              temp_params_dict[param] = mean_field_parameter
          else:
              prior_weight = yield trainable_state_util.Parameter(
                  init_fn=lambda: tf.fill(  # pylint: disable=g-long-lambda
                      dims=param_shape,
                      value=tf.cast(initial_prior_weight, param_dtype),
                  ),
                  name="prior_weight_{}_{}".format(
                      _get_name(dist), param
                  ),
                  constraining_bijector=sigmoid.Sigmoid(),
              )
              temp_params_dict[param] = prior_weight * prior_value + (
                  (1.0 - prior_weight) * mean_field_parameter
              )
      # pylint: enable=cell-var-from-loop
  
      return type(dist)(**temp_params_dict)
  
imports: import collections
  import copy
  import functools
  import inspect
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.bijectors import restructure
  from tensorflow_probability.python.bijectors import scale
  from tensorflow_probability.python.bijectors import shift
  from tensorflow_probability.python.bijectors import sigmoid
  from tensorflow_probability.python.distributions import beta
  from tensorflow_probability.python.distributions import chi2
  from tensorflow_probability.python.distributions import exponential
  from tensorflow_probability.python.distributions import gamma
  from tensorflow_probability.python.distributions import half_normal
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_auto_batched
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import markov_chain
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.distributions import truncated_normal
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_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
label_desc: Use image node_id 10 for calling a global function with example usage: _asvi_convex_update_for_base_distribution(dist, mean_field, initial_prior_weight, sample_shape) and returns: type