flax

Forward experimental checkpoint policy option in lift.remat

#1560

Merged