xla

dd52ee5f - Allow the loop functions to save device specific data (ie, optimizers, in case of them carrying LR schedules) across training loop.