Improve documentation

[AlexeyG]

Let's use this issue to collect all the places in which the documentation should be improved. If you think there's something missing from the list, please comment and we'll update the list.

As we work through this list, we'll update the description with links to PRs addressing each of the items.

• nn.Collection
• Module.apply
• Module.call
• Module should get a docstring
• nn.stateful
• nn.stochastic
• Module.get_param (explain when it should be used)

[avital]

Some of these documentation tasks may be best done as part some broader rethinking of core APIs, such as in #208

[avital]

(And assessing how simple it is to write docstrings for new APIs will help make sure the APIs have a clear mental model)

[lisa-1010]

https://github.com/google/flax/blob/master/flax/training/common_utils.py could benefit from docstrings, especially since the methods in it are used in many flax examples. Thanks!

[avital]

I'll close this issue because it's related to the old flax.nn API and we've now switched to Linen.

@lisa-1010 I filed your comment about common_utils in a separate issue: #729