xla

31919d54 - Add xla autocast support, update autocast APIs in checkpointing (#8523)