xla

15fc0f1c - [FSDPv2] Shard on the maximal dim of weights (#7134)

[FSDPv2] Shard on the maximal dim of weights (#7134) Summary: This pull request makes FSDPv2 to shard on the maximal dim of weights instead of the 0th dim. Test Plan: XLA_USE_SPMD=1 PJRT_DEVICE=TPU python test/spmd/test_fsdp_v2.py