Can device_put used in pmap?

[ZYHowell]

Hi all, I'm wondering if device_put can be used in pmap. In my understand, when two tensors are committed on different devices, they are no longer able to compute. For instance: (default backend is gpu)

import jax
x = jax.device_put(1, jax.devices()[0])
y = jax.device_put(1, jax.local_devices(backend='cpu')[0])
x + y # ValueError: primitive arguments must be colocated on the same device, got gpu:0, cpu:0

However, a likewise case in pmap confused me: (there are 4 gpu devices)

import jax
from jax import pmap, partial
import jax.numpy as np

cpu_device = jax.local_devices(backend='cpu')[0]

@partial(pmap, axis_name="i", in_axes=(0, None))
def mymul(v, W): 
    v_cpu = jax.device_put(v, cpu_device)
    return W.dot(v_cpu)
v = np.ones((4, 3))
W = np.ones((4, 3))
mymul(v, W)

There is no error reported!

So I doubt if device_put really work in pmap? If so, why the code above works? If not, is there a way to transfer a tensor from each device to the host or a certain device in pmap?

[ZYHowell]

Well, I'll explain it myself. The translation rule of device_put is actually lambda c, x, device=None : x in jax/interpreters/xla.py#L1342, so once jit or pmap(which also applies jit) is applied to a function, the device_put inside does not work.

[skye]

Related: #6040