code: def pfor(fn, n):
      if JAX_MODE:
          import jax  # pylint: disable=g-import-not-at-top
  
          return jax.vmap(fn)(np.arange(n))
      outs = [fn(i) for i in range(n)]
      flat_outs = [nest.flatten(o) for o in outs]
      return nest.pack_sequence_as(
          outs[0], [np.array(o) for o in zip(*flat_outs)]
      )
  
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 6 for calling a global function with example usage: pfor(fn, n) and returns: nest, jax