xla

b357d87d - switch to pinning layout by default in FSDP to avoid data corruption in PJRT and be consistent with xm.all_reduce (#4359)