xla

Add optimization barrier to enable the gradient checkpointing

#3482

Merged