flax

[docs] Exposed more helper functions/classes in state.rst

#5037

Merged