xla

28b9b0f8 - Add xla autocast support, update autocast APIs in grad checkpointing (#8527)