code: def map_fn(
      fn,
      elems,
      dtype=None,
      parallel_iterations=None,
      back_prop=True,
      swap_memory=False,
      infer_shape=True,
      name=None,
      fn_output_signature=None,
  ):
      """Transforms `elems` by applying `fn` to each element unstacked on axis 0.
  
      See also `tf.scan`.
  
      `map_fn` unstacks `elems` on axis 0 to obtain a sequence of elements;
      calls `fn` to transform each element; and then stacks the transformed
      values back together.
  
      #### Mapping functions with single-Tensor inputs and outputs
  
      If `elems` is a single tensor and `fn`'s signature is `tf.Tensor->tf.Tensor`,
      then `map_fn(fn, elems)` is equivalent to
      `tf.stack([fn(elem) for elem in tf.unstack(elems)])`.  E.g.:
  
      >>> tf.map_fn(fn=lambda t: tf.range(t, t + 3), elems=tf.constant([3, 5, 2]))
      <tf.Tensor: shape=(3, 3), dtype=int32, numpy=
        array([[3, 4, 5],
               [5, 6, 7],
               [2, 3, 4]], dtype=int32)>
  
      `map_fn(fn, elems).shape = [elems.shape[0]] + fn(elems[0]).shape`.
  
      #### Mapping functions with multi-arity inputs and outputs
  
      `map_fn` also supports functions with multi-arity inputs and outputs:
  
      * If `elems` is a tuple (or nested structure) of tensors, then those tensors
        must all have the same outer-dimension size (`num_elems`); and `fn` is
        used to transform each tuple (or structure) of corresponding slices from
        `elems`.  E.g., if `elems` is a tuple `(t1, t2, t3)`, then `fn` is used to
        transform each tuple of slices `(t1[i], t2[i], t3[i])`
        (where `0 <= i < num_elems`).
  
      * If `fn` returns a tuple (or nested structure) of tensors, then the
        result is formed by stacking corresponding elements from those structures.
  
      #### Specifying `fn`'s output signature
  
      If `fn`'s input and output signatures are different, then the output
      signature must be specified using `fn_output_signature`.  (The input and
      output signatures are differ if their structures, dtypes, or tensor types do
      not match).  E.g.:
  
      >>> tf.map_fn(fn=tf.strings.length,  # input & output have different dtypes
      ...           elems=tf.constant(["hello", "moon"]),
      ...           fn_output_signature=tf.int32)
      <tf.Tensor: shape=(2,), dtype=int32, numpy=array([5, 4], dtype=int32)>
      >>> tf.map_fn(fn=tf.strings.join,  # input & output have different structures
      ...           elems=[tf.constant(['The', 'A']), tf.constant(['Dog', 'Cat'])],
      ...           fn_output_signature=tf.string)
      <tf.Tensor: shape=(2,), dtype=string,
       numpy=array([b'TheDog', b'ACat'], dtype=object)>
  
      `fn_output_signature` can be specified using any of the following:
  
      * A `tf.DType` or `tf.TensorSpec` (to describe a `tf.Tensor`)
      * A `tf.RaggedTensorSpec` (to describe a `tf.RaggedTensor`)
      * A `tf.SparseTensorSpec` (to describe a `tf.sparse.SparseTensor`)
      * A (possibly nested) tuple, list, or dict containing the above types.
  
      #### RaggedTensors
  
      `map_fn` supports `tf.RaggedTensor` inputs and outputs.  In particular:
  
      * If `elems` is a `RaggedTensor`, then `fn` will be called with each
        row of that ragged tensor.
        * If `elems` has only one ragged dimension, then the values passed to
          `fn` will be `tf.Tensor`s.
        * If `elems` has multiple ragged dimensions, then the values passed to
          `fn` will be `tf.RaggedTensor`s with one fewer ragged dimension.
  
      * If the result of `map_fn` should be a `RaggedTensor`, then use a
        `tf.RaggedTensorSpec` to specify `fn_output_signature`.
        * If `fn` returns `tf.Tensor`s with varying sizes, then use a
          `tf.RaggedTensorSpec` with `ragged_rank=0` to combine them into a
          single ragged tensor (which will have ragged_rank=1).
        * If `fn` returns `tf.RaggedTensor`s, then use a `tf.RaggedTensorSpec`
          with the same `ragged_rank`.
  
      >>> # Example: RaggedTensor input
      >>> rt = tf.ragged.constant([[1, 2, 3], [], [4, 5], [6]])
      >>> tf.map_fn(tf.reduce_sum, rt, fn_output_signature=tf.int32)
      <tf.Tensor: shape=(4,), dtype=int32, numpy=array([6, 0, 9, 6], dtype=int32)>
  
      >>> # Example: RaggedTensor output
      >>> elems = tf.constant([3, 5, 0, 2])
      >>> tf.map_fn(tf.range, elems,
      ...           fn_output_signature=tf.RaggedTensorSpec(shape=[None],
      ...                                                   dtype=tf.int32))
      <tf.RaggedTensor [[0, 1, 2], [0, 1, 2, 3, 4], [], [0, 1]]>
  
      Note: `map_fn` should only be used if you need to map a function over the
      *rows* of a `RaggedTensor`.  If you wish to map a function over the
      individual values, then you should use:
  
      * `tf.ragged.map_flat_values(fn, rt)`
        (if fn is expressible as TensorFlow ops)
      * `rt.with_flat_values(map_fn(fn, rt.flat_values))`
        (otherwise)
  
      E.g.:
  
      >>> rt = tf.ragged.constant([[1, 2, 3], [], [4, 5], [6]])
      >>> tf.ragged.map_flat_values(lambda x: x + 2, rt)
      <tf.RaggedTensor [[3, 4, 5], [], [6, 7], [8]]>
  
      #### SparseTensors
  
      `map_fn` supports `tf.sparse.SparseTensor` inputs and outputs.  In particular:
  
      * If `elems` is a `SparseTensor`, then `fn` will be called with each row
        of that sparse tensor. In particular, the value passed to `fn` will be a
        `tf.sparse.SparseTensor` with one fewer dimension than `elems`.
  
      * If the result of `map_fn` should be a `SparseTensor`, then use a
        `tf.SparseTensorSpec` to specify `fn_output_signature`.  The individual
        `SparseTensor`s returned by `fn` will be stacked into a single
        `SparseTensor` with one more dimension.
  
      >>> # Example: SparseTensor input
      >>> st = tf.sparse.SparseTensor([[0, 0], [2, 0], [2, 1]], [2, 3, 4], [4, 4])
      >>> tf.map_fn(tf.sparse.reduce_sum, st, fn_output_signature=tf.int32)
      <tf.Tensor: shape=(4,), dtype=int32, numpy=array([2, 0, 7, 0], dtype=int32)>
  
      >>> # Example: SparseTensor output
      >>> tf.sparse.to_dense(
      ...     tf.map_fn(tf.sparse.eye, tf.constant([2, 3]),
      ...               fn_output_signature=tf.SparseTensorSpec(None, tf.float32)))
      <tf.Tensor: shape=(2, 3, 3), dtype=float32, numpy=
        array([[[1., 0., 0.],
                [0., 1., 0.],
                [0., 0., 0.]],
               [[1., 0., 0.],
                [0., 1., 0.],
                [0., 0., 1.]]], dtype=float32)>
  
      Note: `map_fn` should only be used if you need to map a function over the
      *rows* of a `SparseTensor`.  If you wish to map a function over the nonzero
      values, then you should use:
  
      * If the function is expressible as TensorFlow ops, use:
        ```python
        tf.sparse.SparseTensor(st.indices, fn(st.values), st.dense_shape)
        ```
      * Otherwise, use:
        ```python
        tf.sparse.SparseTensor(st.indices, tf.map_fn(fn, st.values),
                               st.dense_shape)
        ```
  
      #### `map_fn` vs. vectorized operations
  
      `map_fn` will apply the operations used by `fn` to each element of `elems`,
      resulting in `O(elems.shape[0])` total operations.  This is somewhat
      mitigated by the fact that `map_fn` can process elements in parallel.
      However, a transform expressed using `map_fn` is still typically less
      efficient than an equivalent transform expressed using vectorized operations.
  
      `map_fn` should typically only be used if one of the following is true:
  
      * It is difficult or expensive to express the desired transform with
        vectorized operations.
      * `fn` creates large intermediate values, so an equivalent vectorized
        transform would take too much memory.
      * Processing elements in parallel is more efficient than an equivalent
        vectorized transform.
      * Efficiency of the transform is not critical, and using `map_fn` is
        more readable.
  
      E.g., the example given above that maps `fn=lambda t: tf.range(t, t + 3)`
      across `elems` could be rewritten more efficiently using vectorized ops:
  
      >>> elems = tf.constant([3, 5, 2])
      >>> tf.range(3) + tf.expand_dims(elems, 1)
      <tf.Tensor: shape=(3, 3), dtype=int32, numpy=
        array([[3, 4, 5],
               [5, 6, 7],
               [2, 3, 4]], dtype=int32)>
  
      In some cases, `tf.vectorized_map` can be used to automatically convert a
      function to a vectorized equivalent.
  
      #### Eager execution
  
      When executing eagerly, `map_fn` does not execute in parallel even if
      `parallel_iterations` is set to a value > 1. You can still get the
      performance benefits of running a function in parallel by using the
      `tf.function` decorator:
  
      >>> fn=lambda t: tf.range(t, t + 3)
      >>> @tf.function
      ... def func(elems):
      ...   return tf.map_fn(fn, elems, parallel_iterations=3)
      >>> func(tf.constant([3, 5, 2]))
      <tf.Tensor: shape=(3, 3), dtype=int32, numpy=
        array([[3, 4, 5],
               [5, 6, 7],
               [2, 3, 4]], dtype=int32)>
  
  
      Note: if you use the `tf.function` decorator, any non-TensorFlow Python
      code that you may have written in your function won't get executed. See
      `tf.function` for more  details. The recommendation would be to debug without
      `tf.function` but switch to it to get performance benefits of running `map_fn`
      in parallel.
  
      Args:
        fn: The callable to be performed.  It accepts one argument, which will have
          the same (possibly nested) structure as `elems`.  Its output must have the
          same structure as `fn_output_signature` if one is provided; otherwise it
          must have the same structure as `elems`.
        elems: A tensor or (possibly nested) sequence of tensors, each of which will
          be unstacked along their first dimension.  `fn` will be applied to the
          nested sequence of the resulting slices.  `elems` may include ragged and
          sparse tensors. `elems` must consist of at least one tensor.
        dtype: Deprecated: Equivalent to `fn_output_signature`.
        parallel_iterations: (optional) The number of iterations allowed to run in
          parallel. When graph building, the default value is 10. While executing
          eagerly, the default value is set to 1.
        back_prop: (optional) False disables support for back propagation.
        swap_memory: (optional) True enables GPU-CPU memory swapping.
        infer_shape: (optional) False disables tests for consistent output shapes.
        name: (optional) Name prefix for the returned tensors.
        fn_output_signature: The output signature of `fn`. Must be specified if
          `fn`'s input and output signatures are different (i.e., if their
          structures, dtypes, or tensor types do not match).
          `fn_output_signature` can be specified using any of the following:
  
          * A `tf.DType` or `tf.TensorSpec` (to describe a `tf.Tensor`)
          * A `tf.RaggedTensorSpec` (to describe a `tf.RaggedTensor`)
          * A `tf.SparseTensorSpec` (to describe a `tf.sparse.SparseTensor`)
          * A (possibly nested) tuple, list, or dict containing the above types.
  
      Returns:
        A tensor or (possibly nested) sequence of tensors.  Each tensor stacks the
        results of applying `fn` to tensors unstacked from `elems` along the first
        dimension, from first to last.  The result may include ragged and sparse
        tensors.
  
      Raises:
        TypeError: if `fn` is not callable or the structure of the output of
          `fn` and `fn_output_signature` do not match.
        ValueError: if the lengths of the output of `fn` and `fn_output_signature`
          do not match, or if the `elems` does not contain any tensor.
  
      Examples:
  
        >>> elems = np.array([1, 2, 3, 4, 5, 6])
        >>> tf.map_fn(lambda x: x * x, elems)
        <tf.Tensor: shape=(6,), dtype=int64, numpy=array([ 1,  4,  9, 16, 25, 36])>
  
        >>> elems = (np.array([1, 2, 3]), np.array([-1, 1, -1]))
        >>> tf.map_fn(lambda x: x[0] * x[1], elems, fn_output_signature=tf.int64)
        <tf.Tensor: shape=(3,), dtype=int64, numpy=array([-1,  2, -3])>
  
        >>> elems = np.array([1, 2, 3])
        >>> tf.map_fn(lambda x: (x, -x), elems,
        ...          fn_output_signature=(tf.int64, tf.int64))
        (<tf.Tensor: shape=(3,), dtype=int64, numpy=array([1, 2, 3])>,
         <tf.Tensor: shape=(3,), dtype=int64, numpy=array([-1, -2, -3])>)
      """
      # This function uses a `while_loop` to call `fn` on each value of the input
      # tensor(s) (unstacked on dimension 0).  The following sequence of variables
      # are used to transform the input tensor(s) (`elems`) into the output
      # tensor(s) (`result`):
      #
      #   - Preparing and unstacking input values for the while_loop:
      #     - elems: The input tensor(s) to map_fn. May include composite tensors.
      #     - elems_flat: Flattened list of tensors from elems (using nest.flatten)
      #                   May include composite tensors.
      #     - elems_batchable: Concatenation of "batchable tensor lists" for each
      #                        tensor in elems_flat.  This "boxes" composite tensors
      #                        into sliceable tf.Tensor objects.  For more info see:
      #                        TensorSpec._to_batched_tensor_list
      #     - elems_batchable_ta: List of TensorArrays used to unstack each Tensor
      #                           in elems_batchable into elems_value_batchable.
      #
      #   - Calling `fn` on each unstacked value in the body of the while_loop:
      #     - elems_value_batchable: Single unstacked value from elems_batchable.
      #     - elems_value_flat: Single unstacked value from elems_flat,
      #                         constructed from elems_value_batchable (using
      #                         TensorSpec._from_tensor_list).
      #     - elems_value: Single unstacked value from elems (the input to fn).
      #     - result_value: Result of calling `fn(elems_value)`.  May contain
      #                     composite tensors.
      #     - result_value_flat: Flattened list of tensors from result_value.
      #                          May contain composite tensors.
      #     - result_value_batchable: Concatenation of batchable tensor lists for
      #                               each tensor in result_value_flat
      #                               (using TensorSpec._to_tensor_list).
      #
      #   - Collecting and stacking output values from the while_loop:
      #     - result_batchable_ta: List of TensorArrays used to stack each tensor
      #                            ta result_value_batchable into result_batchable.
      #     - result_batchable: Stacked tensors from result_batchable_ta.
      #     - result_flat: Flat list of tensors for the result, constructed from
      #                    results bactchable (using TensorSpec._from_tensor_list).
      #     - result: Structured result value packed from results flat
      #               (using nest.pack_sequence_as).
  
      if fn_output_signature is None:
          fn_output_signature = dtype
  
      if not callable(fn):
          raise TypeError(
              f"The provided function {fn.__name__} is not callable."
              "fn must be callable."
          )
  
      in_graph_mode = not context.executing_eagerly()
      # Set the default number of parallel_iterations depending on graph/eager mode.
      if in_graph_mode and not parallel_iterations:
          parallel_iterations = 10
      elif not in_graph_mode and not parallel_iterations:
          parallel_iterations = 1
      elif not in_graph_mode and parallel_iterations > 1:
          logging.log_first_n(
              logging.WARN,
              "Setting parallel_iterations > 1 has no "
              "effect when executing eagerly. Consider calling map_fn"
              " with tf.function to execute fn in "
              "parallel.",
              1,
          )
          parallel_iterations = 1
  
      # Explicitly read values of ResourceVariables.
      elems = variable_utils.convert_variables_to_tensors(elems)
      # Flatten the input tensors, and get the TypeSpec for each one.
      elems_flat = nest.flatten(elems)
  
      # Check in case this is an empty list
      if len(elems_flat) == 0:
          raise ValueError(
              "elems must be a Tensor or (possibly nested) sequence of Tensors. "
              "Got {}, which does not contain any Tensors.".format(
                  elems
              )
          )
  
      elems_flat_signature = [
          type_spec.type_spec_from_value(e) for e in elems_flat
      ]
      elems_unflatten = lambda x: nest.pack_sequence_as(elems, x)
  
      # Flatten fn's output signature.
      if fn_output_signature is None:
          # If fn_output_signature was not specified, then assume that it matches the
          # input signature.
          result_flat_signature = [
              _most_general_type_spec(
                  e
              )._unbatch()  # pylint: disable=protected-access
              for e in elems_flat
          ]
          result_unflatten = elems_unflatten
      else:
          result_flat_signature = [
              _dtype_to_spec(d)
              for d in nest.flatten(fn_output_signature)
          ]
          result_unflatten = lambda x: nest.pack_sequence_as(
              fn_output_signature, x
          )
  
      with ops.name_scope(name, "map", elems_flat):
          # TODO(akshayka): Remove the in_graph_mode check once caching devices are
          # supported in Eager
          if in_graph_mode:
              # Any get_variable calls in fn will cache the first call locally
              # and not issue repeated network I/O requests for each iteration.
              varscope = vs.get_variable_scope()
              varscope_caching_device_was_none = False
              if varscope.caching_device is None:
                  # TODO(ebrevdo): Change to using colocate_with here and in other
                  # methods.
                  varscope.set_caching_device(lambda op: op.device)
                  varscope_caching_device_was_none = True
  
          elems_flat = [
              ops.convert_to_tensor_or_composite(t, name="elem")
              for t in elems_flat
          ]
  
          # Check that inputs are not scalars.
          first_elem = elems_flat[0]
          if hasattr(first_elem, "shape"):
              elems_static_shape = first_elem.shape
              if (
                  elems_static_shape.ndims is not None
                  and elems_static_shape.ndims < 1
              ):
                  raise ValueError(
                      "Elements in elems must be 1+ dimensional Tensors, not scalars"
                  )
  
          # Box any composite tensors into tensor lists.
          elems_batchable = _elems_flat_to_batchable(elems_flat)
  
          # Find the number of iterations, n.  (may be known statically.)
          n_static = tensor_shape.Dimension(
              tensor_shape.dimension_value(
                  elems_batchable[0]
                  .get_shape()
                  .with_rank_at_least(1)[0]
              )
          )
          for tensor in elems_batchable[1:]:
              n_static.assert_is_compatible_with(
                  tensor_shape.Dimension(
                      tensor_shape.dimension_value(
                          tensor.get_shape().with_rank_at_least(1)[0]
                      )
                  )
              )
          n = n_static.value or array_ops.shape(elems_batchable[0])[0]
  
          # Convert elems to tensor array.
          # TODO(edloper): Should we set infer_shape=False for composite tensors?
          elems_batchable_ta = [
              tensor_array_ops.TensorArray(
                  dtype=t.dtype,
                  size=n,
                  dynamic_size=False,
                  infer_shape=True,
              )
              for t in elems_batchable
          ]
          # Unpack elements
          elems_batchable_ta = [
              ta.unstack(t)
              for (ta, t) in zip(elems_batchable_ta, elems_batchable)
          ]
  
          i = constant_op.constant(0)
  
          # Prepare result tensor array.
          # TODO(edloper): Should we set infer_shape=False for composite tensors?
          result_batchable_tensor_spec = (
              _result_flat_signature_to_batchable_tensor_spec(
                  result_flat_signature
              )
          )
          result_batchable_ta = []
          for spec in result_batchable_tensor_spec:
              result_batchable_ta.append(
                  tensor_array_ops.TensorArray(
                      dtype=spec.dtype,
                      size=n,
                      dynamic_size=False,
                      infer_shape=infer_shape,
                      element_shape=spec.shape,
                  )
              )
  
          def compute(i, tas):
              """The loop body of map_fn.
  
              Args:
                i: the loop counter
                tas: the flat TensorArray accumulator list
  
              Returns:
                (i + 1, tas): the updated counter + updated TensorArrays
  
              Raises:
                TypeError: if fn_output_signature and result_value structure don't match
                ValueType: if fn_output_signature and result_value lengths don't match
              """
              elems_value_batchable = [
                  ta.read(i) for ta in elems_batchable_ta
              ]
              elems_value_flat = _elems_value_batchable_to_flat(
                  elems_value_batchable, elems_flat_signature
              )
              elems_value = elems_unflatten(elems_value_flat)
              ag_ctx = autograph_ctx.control_status_ctx()
              autographed_fn = autograph.tf_convert(fn, ag_ctx)
              result_value = autographed_fn(elems_value)
              nest.assert_same_structure(
                  fn_output_signature or elems, result_value
              )
              result_value_flat = nest.flatten(result_value)
              result_value_batchable = _result_value_flat_to_batchable(
                  result_value_flat, result_flat_signature
              )
              tas = [
                  ta.write(i, value)
                  for (ta, value) in zip(tas, result_value_batchable)
              ]
              return (i + 1, tas)
  
          _, r_a = while_loop.while_loop(
              lambda i, _: i < n,
              compute,
              (i, result_batchable_ta),
              parallel_iterations=parallel_iterations,
              back_prop=back_prop,
              swap_memory=swap_memory,
              maximum_iterations=n,
          )
          result_batchable = [r.stack() for r in r_a]
  
          # Update each output tensor w/ static shape info about the outer dimension.
          for r in result_batchable:
              r.set_shape(
                  tensor_shape.TensorShape(n_static).concatenate(
                      r.get_shape()[1:]
                  )
              )
  
          # TODO(akshayka): Remove the in_graph_mode check once caching devices are
          # supported in Eager
          if in_graph_mode and varscope_caching_device_was_none:
              varscope.set_caching_device(None)
  
          result_flat = _result_batchable_to_flat(
              result_batchable, result_flat_signature, n_static
          )
          result = result_unflatten(result_flat)
          return result
  
imports: import re
  from tensorflow.python.autograph.core import ag_ctx
  from tensorflow.python.autograph.impl import api
  from tensorflow.python.eager import context
  from tensorflow.python.framework import composite_tensor
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_shape
  from tensorflow.python.framework import tensor_spec
  from tensorflow.python.framework import type_spec
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import tensor_array_ops
  from tensorflow.python.ops import variable_scope
  from tensorflow.python.ops import while_loop
  from tensorflow.python.platform import tf_logging
  from tensorflow.python.util import deprecation
  from tensorflow.python.util import nest
  from tensorflow.python.util import variable_utils
  from tensorflow.python.util.tf_export import tf_export
label_desc: Use image node_id 1 for calling a global function with example usage: map_fn(fn, elems, dtype, parallel_iterations, back_prop, swap_memory, infer_shape, name, fn_output_signature) and returns: result, unknown,  tas