xla

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