How to Unshard an Array back to a single device?

[amishra791]

Suppose I had a function that operates on sharded arrays:

import jax
import jax.numpy as jnp

from jax.sharding import PartitionSpec as P, NamedSharding

mesh = jax.make_mesh((2,2), ('x', 'y'))
x = jax.random.normal(jax.random.key(0), (8192, 8192))
w = jax.random.normal(jax.random.key(1), (8192, 8192))
x = jax.device_put(x, NamedSharding(mesh, P('x', 'y')))
w = jax.device_put(w, NamedSharding(mesh, P('x', 'y')))

@jax.jit
def f(x, w):
    y = jnp.dot(x, w)
    return jax.lax.with_sharding_constraint(y, NamedSharding(mesh, P()))

y = f(x, w)

I'd like to make f return an array on a single device, rather than it be replicated or sharded across multiple devices. How can I do that?

[yashk2810]

Outside of the function, you can do arr.addressable_shards and choose any shard you want on any device.

[amishra791]

Thanks! And freeing up memory from the rest of the devices would be a matter of calling del on the rest of the unused shards?

[yashk2810]

Yes, but that would invalidate the current shard on your device too. You probably need to copy the shard and then call arr.delete().