code: def geometric_cooling_schedule(
      cooling_fraction_per_k_iterations, k=1.0
  ):
      """Defines a cooling schedule following a geometric sequence.
  
      This returns a function `f` such that
  
      ```python
      f(iteration) = cooling_fraction_per_k_iterations**(iteration / k)
      ```
  
      Args:
        cooling_fraction_per_k_iterations: float `Tensor` ratio by which the
          original value should be scaled once `k` iterations have been completed.
        k: int `Tensor` number of iterations used to define the schedule.
      Returns:
        f: Python `callable` representing the cooling schedule.
      """
      cooling_fraction_per_k_iterations = tf.convert_to_tensor(
          cooling_fraction_per_k_iterations,
          dtype_hint=tf.float32,
          name="cooling_fraction_per_k_iterations",
      )
      dtype = cooling_fraction_per_k_iterations.dtype
      k = tf.cast(k, dtype=dtype, name="k")
  
      def f(iteration):
          iteration = tf.cast(iteration, dtype=dtype, name="iteration")
          return cooling_fraction_per_k_iterations ** (iteration / k)
  
      return f
  
imports: import collections
  import contextlib
  import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import invert
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.experimental.mcmc.particle_filter import infer_trajectories
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import reparameterization
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 1 for calling a global function with example usage: geometric_cooling_schedule(cooling_fraction_per_k_iterations, k) and returns: f, unknown