flax

3af21539 - Some minor fixes and docstring updates following the last Orbax checkpoint migration PR.