Compress data transfer between GPU and host

[regevs]

I have a computation-intensive function that jax optimizes nicely on a GPU. It turns out most of the actual computation time is dedicated to moving the final result (a jnp.array) back to the host (as a np.array).

Are there any tips as to how speed up this part?

In particular, is there a way to compress the array prior to transferring it back to host?

Thanks!