flax

13b40770 - [nnx] update State and variables docstrings