Operations over ShardedDeviceArrays appear to be very slow

[msdejong]

I'm trying to perform a number of operations over ShardedDeviceArrays, and finding that all of them are orders of magnitudes slower than normal DeviceArrays. I posted #6626 about reshaping, but even for simple operations like indexing, the same problem appears.

For example, I want to index a particular row for each device. Simply indexing that row doesn't work - the resulting array is no longer sharded.

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

import time   
size = 100000
dim = 100
devices = 8

device_array = np.random.rand(dim, size)
device_array = jnp.array(device_array, dtype=jnp.bfloat16)

sharded_array = jax.device_put_replicated(device_array, jax.devices())

test_array = sharded_array[:, 0, :]
print(type(test_array))
<class 'google3.third_party.tensorflow.compiler.xla.python.xla_extension.DeviceArray'>

Using pmap results in a sharded array, but is extremely slow, hundreds of times slower than normal indexing.

ts = %timeit -q  -o -n 10 _ = device_array[0].block_until_ready()
print("Indexing into device array takes {:.3f}ms".format(ts.best * 1000))

ts = %timeit -q  -o -n 10 _ = jax.pmap(lambda x: x[0])(sharded_array)[0].block_until_ready()
print("Indexing into sharded device array takes {:.3f}ms".format(ts.best * 1000))

Indexing into device array takes 1.222ms
Indexing into sharded device array takes 154.576ms

What is the right way to work with arrays on multiple devices?