flax

3e28eeb0 - Use `import orbax.checkpoint as ocp` to avoid the verbosity of using 'orbax.checkpoint` every time.