code: def simple_heuristic_tuning(
      num_steps,
      log_scalings,
      log_accept_prob,
      optimal_accept=0.234,
      target_accept_prob=0.99,
      name=None,
  ):
      """Tune the number of steps and scaling of one mutation.
  
      # TODO(b/152412213): Better explanation of the heuristic used here.
  
      This is a simple heuristic for tuning the number of steps of the next
      mutation, as well as the scaling of a transition kernel (e.g., step size in
      HMC, scale of a Normal proposal in RWMH) using the acceptance probability from
      the previous mutation stage in SMC.
  
      Args:
        num_steps: The initial number of steps for the next mutation, to be tune.
        log_scalings: The log of the scale of the proposal kernel
        log_accept_prob: The log of the acceptance ratio from the last mutation.
        optimal_accept: Optimal acceptance ratio for a Transitional Kernel. Default
          value is 0.234 (Optimal for Random Walk Metropolis kernel).
        target_accept_prob: Target acceptance probability at the end of one mutation
          step. Default value: 0.99
        name: Python `str` name prefixed to Ops created by this function.
          Default value: `None`.
  
      Returns:
        num_steps: The number of steps for the next mutation.
        new_log_scalings: The log of the scale of the proposal kernel for the next
          mutation.
  
      """
      with tf.name_scope(name or "simple_heuristic_tuning"):
          optimal_accept = tf.constant(
              optimal_accept, dtype=log_accept_prob.dtype
          )
          target_accept_prob = tf.constant(
              target_accept_prob, dtype=log_accept_prob.dtype
          )
          log_half_constant = tf.constant(
              np.log(0.5), dtype=log_scalings.dtype
          )
  
          avg_log_scalings = reduce_logmeanexp(log_scalings, axis=0)
          avg_log_accept_prob = reduce_logmeanexp(
              log_accept_prob, axis=0
          )
  
          avg_log_scaling_target = avg_log_scalings + (
              tf.exp(avg_log_accept_prob) - optimal_accept
          )
          new_log_scalings = log_half_constant + log_add_exp(
              avg_log_scaling_target,
              log_scalings + (tf.exp(log_accept_prob) - optimal_accept),
          )
  
          num_particles = ps.shape(log_accept_prob)[-1]
          num_proposed = tf.cast(
              num_particles * num_steps, dtype=avg_log_accept_prob.dtype
          )
          # max(1/num_proposed, average_accept_ratio)
          log_avg_accept = tf.math.maximum(
              -tf.math.log(num_proposed), avg_log_accept_prob
          )
          num_steps = tf.cast(
              tf.math.log1p(-target_accept_prob)
              / log1mexp(log_avg_accept),
              dtype=num_steps.dtype,
          )
          # We choose the number of steps from the batch that takes the longest,
          # hence this is a reduce over all axes.
          max_step_across_batch = tf.reduce_max(num_steps)
          return max_step_across_batch, new_log_scalings
  
imports: import collections
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.mcmc import weighted_resampling
  from tensorflow_probability.python.internal import broadcast_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.math.generic import log1mexp
  from tensorflow_probability.python.math.generic import log_add_exp
  from tensorflow_probability.python.math.generic import reduce_logmeanexp
  from tensorflow_probability.python.mcmc import hmc
  from tensorflow_probability.python.mcmc import random_walk_metropolis
  from tensorflow_probability.python.mcmc import transformed_kernel
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 7 for calling a global function with example usage: simple_heuristic_tuning(num_steps, log_scalings, log_accept_prob, optimal_accept, target_accept_prob, name) and returns: max_step_across_batch,  new_log_scalings