T Coxon

I've updated the import to only require `orbax.checkpointing`

Added orbax-checkpoint as a top level import and a package wide dependency. Still have `flax.training` as a non top level import though.