jax

7d3c63ed - [sharding_in_types] Add more reshape sharding support

[sharding_in_types] Add more reshape sharding support * Allow merging and splitting only if major most dim is sharded since that involves no data movement. This only happens if `dimensions` is None i.e. if the input array is in **row-major order**. * Merging: If **only** the major most dim is sharded of the merge block then that sharding is propagated to the merge block output * Splitting: If the dimension being split is sharded, then the sharding is propagated to the major most dimension post split only if the spec divides the new shape exactly. PiperOrigin-RevId: 730291595