xla

9054e5d6 - Now that ConfigureDistributedTPU releases all the TPU memory, we do not need to tidy up on exit and we can simplify the code.