xla

1d6558a3 - Backport Fix preserve_rng_state for activation checkpointing (#4690) (#6788)