xla

Fix preserve_rng_state for activation checkpointing

#4690

Merged