xla

b5e6b743 - fix spmd module loading with fsdpv2