xla

5710a83c - Add GKE support and various usability improvements in CheckpointManager (#5770)

Add GKE support and various usability improvements in CheckpointManager (#5770) * Add GKE support and various usability improvements in CheckpointManager * Bug fix for async checkpointing fully sharded state dicts