def _pjit(self, f, memory_key=None, donate_argnums=()):
    """JIT f if 1 device is available and pmap if more are available."""
    should_memoize = self._jit_memory is not None and memory_key is not None
    if (should_memoize and memory_key in self._jit_memory):
      logging.info('Found JITed function in memory for: %s', memory_key)
      return self._jit_memory[memory_key]
    if self._n_devices == 1:
      res = fastmath.jit(f, donate_argnums=donate_argnums)
    else:
      res = fastmath.pmap(f, axis_name='batch', donate_argnums=donate_argnums)
    if should_memoize:
      self._jit_memory[memory_key] = res
    return res
JIT f if 1 device is available and pmap if more are available.
trax/optimizers/trainer.py