def _accelerate(f, n_devices):
  """Returns an accelerated version of ``f`` running on ``n_devices``."""
  if n_devices == 0:  # no accelerators - run on CPU
    return fastmath.jit(f, device=jax.local_devices(backend='cpu')[0])

  if n_devices == 1:
    return fastmath.jit(f)

  return fastmath.pmap(f, axis_name='batch')
Returns an accelerated version of ``f`` running on ``n_devices``.
trax/layers/acceleration.py