def on_accelerator(x):
  """Puts ``x`` in (single) accelerator memory in JAX."""
  try:
    accelerator_devices = jax.devices('gpu')
  except RuntimeError:
    try:
      accelerator_devices = jax.devices('tpu')
    except RuntimeError:
      accelerator_devices = []
  if not accelerator_devices:
    return x
  if len(accelerator_devices) != 1:
    return x
  return jax.device_put(x, accelerator_devices[0])
Puts ``x`` in (single) accelerator memory in JAX.
trax/layers/acceleration.py