Viewing jax.Array sharding status in debugger

[imoneoi]

Before the jax.Array migration, when evaluating a variable in the Python debugger, it will be either a ShardedDeviceArray or a DeviceArray, indicating its sharding status. It would show only Array now, and needs an additional evaluation of x.sharding.device_set to determine the sharding status. Consider adding an indicator of sharding status in __repr__ beside the array data?

[sharadmv]

@yashk2810 what do you think?

BTW: you can also use tools like jax.debug.inspect_array_sharding and jax.debug.visualize_array_sharding.

[yashk2810]

Yeah sure, we can do that but it might distract you away from the actual data and things that you really want to see.

You can also just look at x.sharding or use what sharad suggested above.

I am open to adding sharding in the repr. I can prototype something and see how it looks.

[imoneoi]

I think __repr__ like the old behavior of ShardedDeviceArray([...]) looks fine. It gives hint of sharding (single device / multiple device / multiple processes) while inducing minimal distraction.

Additionally, is it difficult to wrap Array that also allows the usage of isinstance(ShardedDeviceArray, x)?

[yashk2810]

is it difficult to wrap Array that also allows the usage of isinstance(ShardedDeviceArray, x)?

It is not possible. ShardedDeviceArray is deprecated. jax.Array is the only type that exists in JAX and you can check it's sharding via .sharding attribute. See https://jax.readthedocs.io/en/latest/jax_array_migration.html#jax-array-migration

To expose sharding information in the repr, we would have to write docs about all kinds of sharding and what it means. I can take a shot at that and then look into exposing sharding on Array's repr