Mixing CPU and GPU data and code

[che-shr-cat]

Hello!

Could you please explain what happens when you JIT-compile code for GPU and send CPU-committed data (or vice versa) to the compiled function? Is the data gets copied to the corresponding device in this case?

JAX does not allow operations between data committed to different devices but allows using data committed to one device in the function compiled for another device. Is it done intentionally? What's the reason behind this?

Example code:

import jax
import jax.numpy as jnp

def selu(x, 
         alpha=1.6732632423543772848170429916717, 
         scale=1.0507009873554804934193349852946):
  '''Scaled exponential linear unit activation function.'''
  return scale * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)

x = jax.random.normal(jax.random.PRNGKey(42), (1_000_000,))

selu_jit_cpu = jax.jit(selu, backend='cpu')
selu_jit_gpu = jax.jit(selu, backend='gpu')

x_cpu = jax.device_put(x, jax.devices('cpu')[0])
x_gpu = jax.device_put(x, jax.devices('gpu')[0])

z = selu_jit_cpu(x_gpu)  # OK

z = selu_jit_gpu(x_cpu)  # OK

x_cpu + x_gpu  # ValueError: primitive arguments must be colocated on the same device (C++ jax.jit). Arguments are on devices: gpu:0 and TFRT_CPU_0

jax.__version__  # 0.3.17