xla

d7d04792 - add `auto_wrap_policy` into XLA FSDP for automatic wrapping (#4318)

add `auto_wrap_policy` into XLA FSDP for automatic wrapping (#4318) * add `auto_wrap_policy` into XLA FSDP for automatic wrapping * address review comments * default to pin_layout=True in MNIST and ImageNet FSDP