flax

improve nn.checkpoint docs

#2837

Merged