code: def _get_step_size(initial_transformed_position, log_prob_fn):
      """Heuristic for initializing step size.
  
      If we (over) optimistically assume good scaling, 1 / sum(event_dims)**0.25
      will be near the optimal step size. We further scale that downwards. See
      Langmore, Ian, Michael Dikovsky, Scott Geraedts, Peter Norgaard, and Rob Von
      Behren. 2019. “A Condition Number for Hamiltonian Monte Carlo.” arXiv
      [stat.CO]. arXiv. http://arxiv.org/abs/1905.09813.
  
      Args:
        initial_transformed_position: Iterable of arguments to log_prob_fn, in order
          to get a dtype and find a heuristic for an initial step size. We assume
          the Tensor has been flattened so that number of event dimensions is the
          last one.
        log_prob_fn: Target log probability function.
  
      Returns:
        Scalar float of the same dtype as log_prob_fn.
      """
      # TODO(b/187658871): Update this code after internal kernels can support it.
      dtype = log_prob_fn(*initial_transformed_position).dtype
      return (
          0.5
          * sum(
              [
                  tf.cast(ps.shape(state_part)[-1], dtype)
                  for state_part in initial_transformed_position
              ]
          )
          ** -0.25
      )
  
imports: import collections
  import functools
  import warnings
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import chain
  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.experimental.mcmc import diagonal_mass_matrix_adaptation
  from tensorflow_probability.python.experimental.mcmc import initialization
  from tensorflow_probability.python.experimental.mcmc import preconditioned_hmc
  from tensorflow_probability.python.experimental.mcmc import preconditioned_nuts
  from tensorflow_probability.python.experimental.mcmc import preconditioning_utils
  from tensorflow_probability.python.experimental.mcmc import sharded
  from tensorflow_probability.python.experimental.stats import sample_stats
  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 tensorshape_util
  from tensorflow_probability.python.internal import unnest
  from tensorflow_probability.python.math import generic
  from tensorflow_probability.python.mcmc import dual_averaging_step_size_adaptation
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc import sample
  from tensorflow_probability.python.mcmc.internal import util
  from tensorflow.python.ops import control_flow_util
  from tensorflow.python.util import nest
label_desc: Use image node_id 11 for calling a global function with example usage: _get_step_size(initial_transformed_position, log_prob_fn) and returns: unknown