Allow on_host/2 to execute certain operations on host

[josevalim]

[josevalim]

For CPU offloading during training: https://huggingface.co/docs/transformers/main_classes/deepspeed

[josevalim]

Btw, we would implement this with infeed, but infeed is not supported in neither iree or cuda pjrt plugins at the moment. There may be a chance this needs to be implemented in a layer above (such as Axon).