xla

Backport Fix preserve_rng_state for activation checkpointing (#4690)

#6788

Merged