xla

dd726d3f - [SPMD] Fix test_spmd_full_to_shard_shape for v2 (#6927)

[SPMD] Fix test_spmd_full_to_shard_shape for v2 (#6927) Summary: This pull request fixes the test case: test_spmd_full_to_shard_shape. Test Plan: PJRT_DEVICE=TPU python test/spmd/test_xla_sharding.py -v -k test_spmd_full_to_shard_shape