xla

6a133b87 - allow FSDP wrapping and sharding over modules on CPU devices (#3992)

allow FSDP wrapping and sharding over modules on CPU devices (#3992) * allow wrapping CPU modules with XLA FSDP and directly sharding params on CPU in this case * update docs