jax

d7c22eb1 - #sdy Do not close any partially sharded dimensions if using auto axes in a shard_map.

#sdy Do not close any partially sharded dimensions if using auto axes in a shard_map. This change reverts cl/731724837 ([github link](https://github.com/jax-ml/jax/commit/4997e45743e3b243ef153674a11a826843ab37b0)), which is a temporary solution to solve the inconsistent padding on the boundary of manual computation. Now that we have a better solution cl/746600070 ([github link](https://github.com/openxla/xla/pull/25080)), we revert this temporary solution. We still keep the added `test_partially_sharded_dim_with_auto` to verify the correctness. Reverts 4997e45743e3b243ef153674a11a826843ab37b0 PiperOrigin-RevId: 749123936