jax

5ef7deb7 - [jax:custom_partitioning] Allow factors for non-batching dimensions to

[jax:custom_partitioning] Allow factors for non-batching dimensions to correspond to different sizes and choose the smallest size to use in the sharing rule. PiperOrigin-RevId: 786798901