Device used by explicit staging

[Gattocrucco]

In the documentation on staging I could not find for what device the function is compiled for.

After a bit of experimentation, I think the criterion is that it's compiled for the device used in the arguments to lower. If the compiled function is called with arguments on another device, they are silently moved to the device assigned to the compiled function, and the results are on the device of the compiled function as well.

Is my guess correct?

Example:

import jax
f = jax.jit(lambda x: x)
xcpu = jax.device_put(0, jax.devices('cpu')[0])
xgpu = jax.device_put(0, jax.devices('gpu')[0])
fcpu = f.lower(xcpu).compile()
fgpu = f.lower(xgpu).compile()
print(fcpu(xcpu).devices()) # -> cpu
print(fcpu(xgpu).devices()) # -> cpu
print(fgpu(xcpu).devices()) # -> gpu
print(fgpu(xgpu).devices()) # -> gpu

{CpuDevice(id=0)}
{CpuDevice(id=0)}
{cuda(id=0)}
{cuda(id=0)}