code: def smart_for_loop(
      loop_num_iter,
      body_fn,
      initial_loop_vars,
      parallel_iterations=10,
      unroll_threshold=1,
      name=None,
  ):
      """Construct a for loop, preferring a python loop if `n` is statically known.
  
      Given `loop_num_iter` and `body_fn`, return an op corresponding to executing
      `body_fn` `loop_num_iter` times, feeding previous outputs of `body_fn` into
      the next iteration.
  
      If `loop_num_iter` is statically known, the op is constructed via python for
      loop, and otherwise a `tf.while_loop` is used.
  
      Args:
        loop_num_iter: `Integer` `Tensor` representing the number of loop
          iterations.
        body_fn: Callable to be executed `loop_num_iter` times.
        initial_loop_vars: Listlike object of `Tensors` to be passed in to
          `body_fn`'s first execution.
        parallel_iterations: The number of iterations allowed to run in parallel.
          It must be a positive integer. See `tf.while_loop` for more details.
          Default value: `10`.
        unroll_threshold: Integer denoting the maximum number of iterations to
          unroll, if possible. If `loop_num_iter > unroll_threshold` a
          `tf.while_loop` will always be used, even if `loop_num_iter` is
          statically known.
          Default value: `1`.
        name: Python `str` name prefixed to Ops created by this function.
          Default value: `None` (i.e., "smart_for_loop").
      Returns:
        result: `Tensor` representing applying `body_fn` iteratively `n` times.
      """
      with tf.name_scope(name or "smart_for_loop"):
          loop_num_iter_ = tf.get_static_value(loop_num_iter)
          if (
              loop_num_iter_ is None
              or tf.executing_eagerly()
              # large values for loop_num_iter_ will cause ridiculously slow
              # graph compilation time (GitHub issue #1033)
              or loop_num_iter_ > unroll_threshold
              or control_flow_util.GraphOrParentsInXlaContext(
                  tf1.get_default_graph()
              )
          ):
              # Cast to int32 to run the comparison against i in host memory,
              # where while/LoopCond needs it.
              loop_num_iter = tf.cast(loop_num_iter, dtype=tf.int32)
              return tf.while_loop(
                  cond=lambda i, *args: i < loop_num_iter,
                  body=lambda i, *args: [i + 1] + list(body_fn(*args)),
                  loop_vars=[np.int32(0)] + initial_loop_vars,
                  parallel_iterations=parallel_iterations,
              )[1:]
          result = initial_loop_vars
          for _ in range(loop_num_iter_):
              result = body_fn(*result)
          return result
  
file_path: probability/tensorflow_probability/python/internal/loop_util.py
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow.python.ops import control_flow_util
label_desc: Use image node_id 4 for calling a global function with example usage: smart_for_loop(loop_num_iter, body_fn, initial_loop_vars, parallel_iterations, unroll_threshold, name) and returns: result, tf