Switching computation between CPU and GPU inside a JIT-ed function

[fbartolic]

I have one small function, say fn_cpu, which is not well suited to run on a GPU. I used a decorator @partial(jit, backend="cpu") to force it to execute on a CPU. I then use this function in another function decorated with @jit without the backend option.

I naively assumed that if I specify the GPU backend globally with config.update("jax_platform_name", "gpu") when fn_cpu is called the DeviceArray is automatically going to be copied from GPU memory to CPU memory, fn_cpu is going to execute on a cpu, and the result will be automatically copied back into GPU memory. I guess it doesn't work like that because I'm getting an error:

ValueError: Outer-jit backend specification cuda must match explicit inner-jit backend specification cpu.

Is there a way to get around this problem? Do I have to manually transfer the arrays using jax.device_put or something similar?

Thank you!

[YouJiacheng]

Check this: https://jax.readthedocs.io/en/latest/jax.experimental.host_callback.html

[YigitElma]

Since the above function is now deprecated, what is the optimum way to execute part of the jited function on the CPU and the other part on the GPU? I would like to be able to call a previously jited function that runs specifically on a CPU, from another jited function on GPU, but probably the outermost jit compiles everything such that the previous compilation is ditched.