flax

0cb9e99a - Add checkpointing documentation and fix clash of jupyter and tailing_whitespace