flax

41c9304c - Make Flax checkpointing use `sharding` instead of extracting mesh and pspecs. Also improved some warning messages.