transform ShardedDeviceArrays to numpy Array appear to be very slow

[iopcaslyz]

When I transform ShardedDeviceArray to numpy Array for the in-place operation, the process appear to be very slow.
I use the different method to achieve transformation:

  import jax.numpy as jnp
  import numpy as np
  import jax
  out_0 = jax.pmap(lambda x: x ** 2)(np.arange(114*114*10000).reshape(2,5000,114,114))
  %time a=jax.device_get(out_0)

CPU times: user 233 ms, sys: 303 ms, total: 536 ms
Wall time: 536 ms

out_0 = jax.pmap(lambda x: x ** 2)(np.arange(114*114*10000).reshape(2,5000,114,114))
%time a=np.array(out_0)

CPU times: user 317 ms, sys: 393 ms, total: 710 ms
Wall time: 710 ms