What is OpShardingSharding?

[functionstackx]

In regards to jax.sharding, what is OpShardingSharding?

cc: @yashk2810

[yashk2810]

Hey, good question!

OpShardingSharding is the base XLACompatibleSharding that every subclass of XLACompatibleSharding eventually lower to inside jax.

The main thing is it consists of an OpSharding. OpSharding is the representation that XLA understands. It is basically documented here: https://cs.opensource.google/tensorflow/tensorflow/+/master:tensorflow/compiler/xla/xla_data.proto;l=715-765?q=xla_data.proto&ss=tensorflow%2Ftensorflow though it may not be very clear on how things work.

To read the OpSharding, you look at the type, tile_assignment_dimensions and tile_assignment_devices. tile_assignment_dimension will tell you how many ways a dimension is sharded or replicated.

If you see something like [4, 2]{0, 1, 2, 3} then this means that the first dimension is sharded 4 ways and 2 dimension is sharded 2 ways and there are 4 devices present {0, 1, 2, 3}.

If you see: [4, 1, 2]{0, 1, 2, 3} last_dim_replicated, then the first dimension is sharded 4 ways, 2nd dimension is sharded 1 way. But the last_dim_replicated here means that the 2nd dimension is actually replicated 2 ways.

This is a simple function we use to figure this information out: https://github.com/google/jax/blob/main/jax/interpreters/pxla.py#L211-L223.

[functionstackx]

thanks for the detailed answer :)

[sharadmv]

Note that OpShardingShardings pop up as a result of XLA determining intermediate shardings inside of a jitted function. You can provide inputs that are sharded with NamedShardings and PositionalShardings but the intermediates are up to XLA (and XLA uses OpShardings as its internal representation).

[yashk2810]

Also Jax can choose to return OpShardingShardings on the boundary too i.e. input and outputs