code: def _scan(  # pylint: disable=unused-argument
      fn,
      elems,
      initializer=None,
      parallel_iterations=10,
      back_prop=True,
      swap_memory=False,
      infer_shape=True,
      reverse=False,
      name=None,
  ):
      """Scan implementation."""
  
      if reverse:
          elems = nest.map_structure(lambda x: x[::-1], elems)
  
      if initializer is None:
          initializer = nest.map_structure(
              lambda x: x[0], elems, expand_composites=True
          )
          elems = nest.map_structure(
              lambda x: x[1:], elems, expand_composites=True
          )
          prepend = initializer
      else:
          prepend = None
  
      if JAX_MODE:
          from jax import lax  # pylint: disable=g-import-not-at-top
  
          def scan_body(arg, x):
              arg = fn(arg, x)
              return arg, arg
  
          _, out = lax.scan(scan_body, initializer, elems)
      else:
          length = len(nest.flatten(elems)[0])
          arg = initializer
          out = []
          for i in range(length):
              arg = fn(
                  arg, nest.map_structure(lambda x: x[i], elems)
              )  # pylint: disable=cell-var-from-loop
              out.append(arg)
          out = nest.map_structure(lambda *x: np.stack(x, axis=0), *out)
  
      if prepend is not None:
          out = nest.map_structure(
              lambda p, o: np.concatenate([p[np.newaxis], o], axis=0),
              prepend,
              out,
          )
  
      ordering = (lambda x: x[::-1]) if reverse else (lambda x: x)
      return nest.map_structure(ordering, out, expand_composites=True)
  
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 5 for calling a global function with example usage: _scan(fn, elems, initializer, parallel_iterations, back_prop, swap_memory, infer_shape, reverse, name) and returns: nest, arg,  arg