code: def trace(state, fn, num_steps, unroll, parallel_iterations=10):
      """TF implementation of `trace` operator, without the calling convention."""
      if (
          tf.config.experimental_functions_run_eagerly()
          or tf.executing_eagerly()
      ):
          state, first_untraced, first_traced = fn(state)
          arrays = tf.nest.map_structure(
              lambda v: tf.TensorArray(  # pylint: disable=g-long-lambda
                  v.dtype, size=num_steps, element_shape=v.shape
              ).write(0, v),
              first_traced,
          )
          start_idx = 1
      else:
          # We need the shapes and dtypes of the outputs of `fn` function to create
          # the `TensorArray`s etc., we can get it by pre-compiling the wrapper
          # function.
          input_spec = tf.nest.map_structure(
              tf.TensorSpec.from_tensor, state
          )
          fn, (_, untraced_spec, traced_spec) = _eval_shape(
              fn, input_spec
          )
  
          arrays = tf.nest.map_structure(
              lambda spec: tf.TensorArray(  # pylint: disable=g-long-lambda
                  spec.dtype, size=num_steps, element_shape=spec.shape
              ),
              traced_spec,
          )
          first_untraced = tf.nest.map_structure(
              lambda spec: tf.zeros(spec.shape, spec.dtype),
              untraced_spec,
          )
          start_idx = 0
  
      def body(i, state, _, arrays):
          state, untraced, traced = fn(state)
          arrays = tf.nest.map_structure(
              lambda a, e: a.write(i, e), arrays, traced
          )
          return i + 1, state, untraced, arrays
  
      def cond(i, *_):
          return i < num_steps
  
      static_num_steps = tf.get_static_value(num_steps)
      loop_vars = (start_idx, state, first_untraced, arrays)
  
      if unroll:
          if static_num_steps is None:
              raise ValueError(
                  "Cannot unroll when `num_steps` is not statically known."
              )
          # TODO(siege): Investigate if using lists instead of TensorArray's is faster
          # (like is done in the JAX backend).
          for _ in range(start_idx, static_num_steps):
              loop_vars = body(*loop_vars)
          _, state, untraced, arrays = loop_vars
      else:
          if static_num_steps is None:
              maximum_iterations = None
          else:
              maximum_iterations = static_num_steps - start_idx
          _, state, untraced, arrays = tf.while_loop(
              cond=cond,
              body=body,
              loop_vars=loop_vars,
              parallel_iterations=parallel_iterations,
              maximum_iterations=maximum_iterations,
          )
  
      traced = tf.nest.map_structure(lambda a: a.stack(), arrays)
  
      def _merge_static_length(x):
          x.set_shape(
              tf.TensorShape(static_num_steps).concatenate(x.shape[1:])
          )
          return x
  
      traced = tf.nest.map_structure(_merge_static_length, traced)
  
      return state, untraced, traced
  
imports: import functools
  import numpy
  import six
  import tensorflow.compat.v2
  from tensorflow.python.util import nest
label_desc: Use image node_id 16 for calling a global function with example usage: trace(state, fn, num_steps, unroll, parallel_iterations) and returns: state,  untraced,  traced, unknown,  state,  untraced,  arrays, unknown, x