flax
updated NNX nn docstrings
#3972
Merged

Loading