code: def _map_fn(  # pylint: disable=unused-argument
      fn,
      elems,
      dtype=None,
      parallel_iterations=None,
      back_prop=True,
      swap_memory=False,
      infer_shape=True,
      name=None,
      fn_output_signature=None,
  ):
      """Numpy implementation of tf.map_fn."""
      if fn_output_signature is not None and nest.is_nested(
          fn_output_signature
      ):
          # If fn returns a tuple, then map_fn returns a tuple as well; and similarly
          # for lists and more complex nestings.  We do not support this behavior at
          # this time, so we raise an error explicitly instead of silently doing the
          # wrong thing.
          raise NotImplementedError
      if JAX_MODE:
          from jax import (
              tree_util,
          )  # pylint: disable=g-import-not-at-top
  
          elems_flat, in_tree = tree_util.tree_flatten(elems)
          elems_zipped = zip(*elems_flat)
  
          def func(flat_args):
              unflat_args = tree_util.tree_unflatten(in_tree, flat_args)
              return fn(unflat_args)
  
          return np.stack([func(x) for x in elems_zipped])
  
      if isinstance(elems, np.ndarray):
          return np.array([fn(x) for x in elems])
  
      # In the NumPy backend, we do not yet support map_fn over lists, tuples, or
      # other structures.
      raise NotImplementedError
  
file_path: probability/tensorflow_probability/python/internal/backend/numpy/functional_ops.py
imports: import numpy
  from tensorflow_probability.python.internal.backend.numpy import _utils
  from tensorflow_probability.python.internal.backend.numpy import nest
label_desc: Use image node_id 3 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: np, np, fn