code: def _vectorized_map(fn, elems):
      """Numpy implementation of tf.vectorized_map."""
      if JAX_MODE:
          from jax import vmap  # pylint: disable=g-import-not-at-top
  
          return vmap(fn)(elems)
  
      # In the NumPy backend, we don't actually vectorize.
      return _map_fn(fn, elems)
  
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 4 for calling a global function with example usage: _vectorized_map(fn, elems) and returns: _map_fn, vmap