xla

46e62380 - [SPMD] asynchronous data loading with ParallelLoader (#4808)

[SPMD] asynchronous data loading with ParallelLoader (#4808) * Allow input sharding spec to be passed to ParallelLoader * Use async data loading for spmd test scripts