flax

Remove string from checkpointing example and unbreak doctest in head

#3288

Merged