code: def make_simple_step_size_update_policy(
      num_adaptation_steps,
      target_rate=0.75,
      decrement_multiplier=0.01,
      increment_multiplier=0.01,
      step_counter=None,
  ):
      """Create a function implementing a step-size update policy.
  
      The simple policy increases or decreases the `step_size_var` based on the
      average of `exp(minimum(0., log_accept_ratio))`. It is based on
      [Section 4.2 of Andrieu and Thoms (2008)](
      https://people.eecs.berkeley.edu/~jordan/sail/readings/andrieu-thoms.pdf).
  
      The `num_adaptation_steps` argument is set independently of any burnin
      for the overall chain. In general, adaptation prevents the chain from
      reaching a stationary distribution, so obtaining consistent samples requires
      `num_adaptation_steps` be set to a value [somewhat smaller](
      http://andrewgelman.com/2017/12/15/burn-vs-warm-iterative-simulation-algorithms/#comment-627745)
      than the number of burnin steps. However, it may sometimes be helpful to set
      `num_adaptation_steps` to a larger value during development in order to
      inspect the behavior of the chain during adaptation.
  
      Args:
        num_adaptation_steps: Scalar `int` `Tensor` number of initial steps to
          during which to adjust the step size. This may be greater, less than, or
          equal to the number of burnin steps. If `None`, the step size is adapted
          on every step (note this breaks stationarity of the chain!).
        target_rate: Scalar `Tensor` representing desired `accept_ratio`.
          Default value: `0.75` (i.e., [center of asymptotically optimal
          rate](https://arxiv.org/abs/1411.6669)).
        decrement_multiplier: `Tensor` representing amount to downscale current
          `step_size`.
          Default value: `0.01`.
        increment_multiplier: `Tensor` representing amount to upscale current
          `step_size`.
          Default value: `0.01`.
        step_counter: Scalar `int` `Variable` specifying the current step. The step
          size is adapted iff `step_counter < num_adaptation_steps`.
          Default value: if `None`, an internal variable
            `step_size_adaptation_step_counter` is created and initialized to `-1`.
  
      Returns:
        step_size_simple_update_fn: Callable that takes args
          `step_size_var, kernel_results` and returns updated step size(s).
      """
      if step_counter is None and num_adaptation_steps is not None:
          step_counter = tf1.get_variable(
              name="step_size_adaptation_step_counter",
              initializer=tf.constant(-1, dtype=tf.int32),
              # Specify the dtype for variable sharing to work correctly
              # (b/120599991).
              dtype=tf.int32,
              trainable=False,
              use_resource=True,
          )
  
      def step_size_simple_update_fn(step_size_var, kernel_results):
          """Updates (list of) `step_size` using a standard adaptive MCMC procedure.
  
          Args:
            step_size_var: (List of) `tf.Variable`s representing the per `state_part`
              HMC `step_size`.
            kernel_results: `collections.namedtuple` containing `Tensor`s
              representing values from most recent call to `one_step`.
  
          Returns:
            step_size_assign: (List of) `Tensor`(s) representing updated
              `step_size_var`(s).
          """
  
          if kernel_results is None:
              if mcmc_util.is_list_like(step_size_var):
                  return [tf.identity(ss) for ss in step_size_var]
              return tf.identity(step_size_var)
          log_n = tf.math.log(
              tf.cast(
                  tf.size(kernel_results.log_accept_ratio),
                  kernel_results.log_accept_ratio.dtype,
              )
          )
          log_mean_accept_ratio = (
              tf.reduce_logsumexp(
                  tf.minimum(kernel_results.log_accept_ratio, 0.0)
              )
              - log_n
          )
          adjustment = tf.where(
              log_mean_accept_ratio
              < tf.cast(
                  tf.math.log(target_rate), log_mean_accept_ratio.dtype
              ),
              -decrement_multiplier / (1.0 + decrement_multiplier),
              increment_multiplier,
          )
  
          def build_assign_op():
              if mcmc_util.is_list_like(step_size_var):
                  return [
                      ss.assign_add(ss * tf.cast(adjustment, ss.dtype))
                      for ss in step_size_var
                  ]
              return step_size_var.assign_add(
                  step_size_var
                  * tf.cast(adjustment, step_size_var.dtype)
              )
  
          if num_adaptation_steps is None:
              return build_assign_op()
          else:
              with tf.control_dependencies(
                  [step_counter.assign_add(1)]
              ):
                  return tf.cond(
                      pred=step_counter < num_adaptation_steps,
                      true_fn=build_assign_op,
                      false_fn=lambda: step_size_var,
                  )
  
      return step_size_simple_update_fn
  
imports: import collections
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import distribute_lib
  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.mcmc import kernel
  from tensorflow_probability.python.mcmc import metropolis_hastings
  from tensorflow_probability.python.mcmc.internal import leapfrog_integrator
  from tensorflow_probability.python.mcmc.internal import util
  from tensorflow.python.util import deprecation
label_desc: Use image node_id 1 for calling a global function with example usage: make_simple_step_size_update_policy(num_adaptation_steps, target_rate, decrement_multiplier, increment_multiplier, step_counter) and returns: step_size_simple_update_fn, tf, step_size_var, build_assign_op, unknown, unknown, tf