What happens when an array in cpu/gpu is passed to a function JITed in gpu/cpu

[muupan]

I noticed that, even if a function is JITed with backend="cpu", I can pass an array in gpu to it, and vice versa. I have questions on this behavior.

• Let me call the device specified at JIT A and the device of the input B. What happens when A and B are different? It seems like returned values are in A, not in B. Is the input copied to A? Does the actual computation of the function happen in A?
• This behavior is unexpected to me. Are there any performance-related concerns when this happens? Is there a way to detect when this happens or prevent it?

import jax
print(jax.__version__, jax.default_backend())

def f(x):
    return x ** 2

fc = jax.jit(f, backend="cpu")
fg = jax.jit(f, backend="gpu")

xc = jax.device_put(jnp.arange(3), device=jax.devices("cpu")[0])
xg = jax.device_put(jnp.arange(3), device=jax.devices("gpu")[0])

# must be ok
fcxc = fc(xc)
print(fcxc, fcxc.device())
fgxg = fg(xg)
print(fgxg, fgxg.device())

# run without error, but what is happening?
fcxg = fc(xg)
print(fcxg, fcxg.device())
fgxc = fg(xc)
print(fgxc, fgxc.device())

0.3.23 gpu
[0 1 4] TFRT_CPU_0
[0 1 4] gpu:0
[0 1 4] TFRT_CPU_0
[0 1 4] gpu:0