How to do fast reshape over shardeddevicearray?

[msdejong]

The same reshape operation on a shardeddevicearray is thousands of times slower than the same operation on an individual device. What is the best way to reshape a sharded array? I have tried both normal reshape and pmap reshape on a v2 TPU (code below), and they are both incredibly slow.

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

msize = 100000
dim = 100
devices = 8

mdevice = np.random.rand(dim, msize)
mdevice = jnp.array(device_memory, dtype=jnp.bfloat16)
msharded = jax.device_put_replicated(mdevice, jax.devices())

ts = %timeit -q  -o -n 10 _ = jnp.reshape(mdevice, (-1)).block_until_ready()
print("Reshaping takes {:.3f}ms".format(ts.best * 1000))

ts = %timeit -q  -o -n 10 _ = jnp.reshape(msharded, (devices, -1)).block_until_ready()
print("Reshaping sharded takes {:.3f}ms".format(ts.best * 1000))

ts = %timeit -q  -o -n 10 _ = jax.pmap(lambda x: jnp.reshape(x, (-1)))(msharded).block_until_ready()
print("Reshaping pmap sharded takes {:.3f}ms".format(ts.best * 1000))