How to localize the distributed jax.device array?

[jaeyoo]

Let's assume a jax.DeviceArray is distributed over multiple cores, but it's small enough to be located at a core. so I want to localize it again. How can it be possible in Jax? is there a helper function to do it? Or, at least I want to localize the data into a few neighbor cores.

[zhangqiaorjc]

The following should work. I don't know if it's the optimal way to do this regarding the number of host copies. cc @jekbradbury

import jax

def f(key):
  key, subkey = jax.random.split(key)
  return jax.random.uniform(subkey, shape=(3,))

prng_key = jax.random.PRNGKey(41)

# each device gets a different prng_key
device_keys = jax.random.split(prng_key, num=jax.local_device_count())
result = jax.pmap(f)(device_keys)
# note that each core produces a different (3,) uniform random values
print(repr(result))

# ShardedDeviceArray([[0.0898279 , 0.21871209, 0.975634  ],
#                    [0.92832124, 0.5143677 , 0.32663155],
#                    [0.5004319 , 0.42006207, 0.5135381 ],
#                    [0.5356785 , 0.63922024, 0.99091077],
#                    [0.8510555 , 0.77760315, 0.60934937],
#                    [0.4079821 , 0.54433346, 0.5355407 ],
#                    [0.8786584 , 0.1137749 , 0.43614888],
#                    [0.53266203, 0.9267961 , 0.26775682]], dtype=float32)

y = result[:]
y

#DeviceArray([[0.0898279 , 0.21871209, 0.975634  ],
#             [0.92832124, 0.5143677 , 0.32663155],
#             [0.5004319 , 0.42006207, 0.5135381 ],
#             [0.5356785 , 0.63922024, 0.99091077],
#             [0.8510555 , 0.77760315, 0.60934937],
#             [0.4079821 , 0.54433346, 0.5355407 ],
#             [0.8786584 , 0.1137749 , 0.43614888],
#             [0.53266203, 0.9267961 , 0.26775682]], dtype=float32)

z = jax.device_put_sharded(jnp.split(y, 2), devices=jax.devices()[:2])
z

#ShardedDeviceArray([[[0.0898279 , 0.21871209, 0.975634  ],
#                     [0.92832124, 0.5143677 , 0.32663155],
#                     [0.5004319 , 0.42006207, 0.5135381 ],
#                     [0.5356785 , 0.63922024, 0.99091077]],
#
#                    [[0.8510555 , 0.77760315, 0.60934937],
#                     [0.4079821 , 0.54433346, 0.5355407 ],
#                     [0.8786584 , 0.1137749 , 0.43614888],
#                     [0.53266203, 0.9267961 , 0.26775682]]],                   dtype=float32)

[zhangqiaorjc]

Note that

y.device()
# TpuDevice(id=0, host=0, coords=(0,0,0), core_on_chip=0)

z.device_buffers[0].device()`
# TpuDevice(id=0, host=0, coords=(0,0,0), core_on_chip=0)

z.device_buffers[1].device()`
# TpuDevice(id=1, host=0, coords=(0,0,0), core_on_chip=1)

[jaeyoo]

Thanks a lot!