code: def _map(
      map_fn: Callable[[Array], jnp.ndarray],
      mode: str,
  ) -> Callable[[jnp.ndarray], jnp.ndarray]:
      """A versatile vmap-like function."""
      if mode == "auto":
          # Current tests show that for our purpose 'map' is better than the other
          # options on all platforms ('gpu', 'tpu', 'cpu'). But this may change in the
          # future.
          mode = "map"
      if mode == "loop":
          return lambda v: jnp.stack([map_fn(e) for e in v], axis=0)
      elif mode == "map":
          return lambda v: jax.lax.map(map_fn, v)
      elif mode == "vmap":
          return jax.vmap(map_fn)
      else:
          raise ValueError("Invalid batching mode.")
  
file_path: scenic/scenic/projects/baselines/deformable_detr/attention.py
imports: import functools
  import math
  from typing import Callable, Sequence, Tuple, Union
  import flax.linen
  import jax
  from jax.core import Shape
  import jax.numpy
  import ml_collections
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: _map(map_fn, mode) and returns: function, function, jax