MPSD (multiple program single data) parallelism?

[f0uriest]

JAX has great support for SIMD/SPMD parallelism through vmap, pmap, xmap, shmap etc, but I'm wondering if the converse is possible.

Basically, I have the following situation:

@jit
@jax.grad
def myfun(x):
    y1 = expensive_fun_on_device1(x)
    y2 = expensive_fun_on_device2(x)
    y3 = expensive_fun_on_device1(x)
    ...
    z = cheapish_fun_on_host_or_anywhere(y1, y2, y3, ...)
    return z

Without jit I can use device_put/device_get to force different calculations to happen on different devices, but I'd like to be able to jit the entire thing (this is library code that the uesr may want to apply arbitrary other transforms to). Also in general y1, y2, y3,... will have different shapes so some of the usual tricks won't work.

Has anyone dealt with this design pattern before or have any ideas?