flax

b97757cf - improve nn.map_variables docstrings