xla

89b81197 - [SPMD] Add FSDP sharding for test_train_spmd_linear_model.py

[SPMD] Add FSDP sharding for test_train_spmd_linear_model.py Summary: This diff adds FSDP sharding for test_train_spmd_linear_model.py. Test Plan: PJRT_DEVICE=TPU XLA_USE_SPMD=1 python test/spmd/test_train_spmd_linear_model.py --sharding fsdp