def on_cpu(x):
  """Puts ``x`` in CPU memory in JAX."""
  if fastmath.is_backend(fastmath.Backend.JAX):
    return jax.device_put(x, jax.local_devices(backend='cpu')[0])
  else:
    return x
Puts ``x`` in CPU memory in JAX.
trax/layers/acceleration.py