xla

fba326e8 - Expose replica_id in XLAShard (#5654)

Expose replica_id in XLAShard (#5654) * Expose shard rank in XLAShard * Improve tests and fix dist chkpt * Change 'rank' to 'replica_id' * Improve documentation