xla

a5692c20 - [SPMD] input sharding for both train & test data loading in imagenet example (#6515)