xla

4aa4dcaf - Implement SPMDSavePlanner to take distributed checkpoints (#5170)

Implement SPMDSavePlanner to take distributed checkpoints (#5170) * Implement SPMDSavePlanner to take distributed checkpoints * Update comment * Duplicate unstable upstream APIs * Duplicate private API narrow_tensor_by_index