How to Manually Inspect Array Shardings?

[amishra791]

Hi, suppose I shard an array like so:

mesh = jax.make_mesh((2, 4), ('x', 'y'))
sharding = jax.sharding.NamedSharding(mesh, P('x', 'y'))
arr_sharded = jax.device_put(arr, sharding)

Is there any way to see what subarray a particular device stores?

[yashk2810]

Yeah, you can use arr_sharded.addressable_shards