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):
  """
  """
  if step_counter is None and num_adaptation_steps is not None:
    step_counter = tf.compat.v1.get_variable(
        name='step_size_adaptation_step_counter',
        initializer=np.array(-1, dtype=np.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(input=kernel_results.log_accept_ratio),
            kernel_results.log_accept_ratio.dtype))
    log_mean_accept_ratio = tf.reduce_logsumexp(
        input_tensor=tf.minimum(kernel_results.log_accept_ratio, 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. + 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