How do I combine jit-ted computations across different devices?

[f2013519]

How would one combine results on jit-ted computations across different devices?

I have a use case where I would like to run some part of the computation on CPU and transfer the result to the GPU to run the remaining compute.

Also in general, not just a CPU, lets say I have a HW accelerator backend that performs some compute. How do I pass that result to an xPU for other compute?

The below code errors saying: "Received incompatible devices for jitted computation. Got argument x of combine with shape float32[] and device ids [0] on platform GPU and argument y of combine with shape float32[] and device ids [0] on platform CPU"

import jax
import jax.numpy as jnp
import numpy as np

cuda_devices = jax.devices('cuda')
cpu_devices = jax.devices('cpu')

gpu_sharding = jax.sharding.SingleDeviceSharding(cuda_devices[0])
cpu_sharding = jax.sharding.SingleDeviceSharding(cpu_devices[0])

x = jnp.array(np.random.normal(size=(4096, 7, 64)))

@jax.jit
def f_gpu(x):
  x = jax.lax.with_sharding_constraint(x, gpu_sharding)
  return jnp.linalg.pinv(x).sum()
z1 = jax.block_until_ready(f_gpu(x))

@jax.jit
def f_cpu(x):
  x = jax.lax.with_sharding_constraint(x, cpu_sharding)
  return jnp.linalg.pinv(x).sum()
z2 = jax.block_until_ready(f_cpu(x))

@jax.jit
def combine(x, y):
    a =  x + y
    return a

#Error 1
combine(z1, z2) 
#Error 2
y = z1 + z2

[f2013519]

I guess you just need to use jax.device_put