Migrate initializers and activation functions to jax.nn

[jekbradbury]

This PR does a few things:

• Adds a simple truncated normal sampler, which I've manually verified does the right thing for inputs in the regime relevant to neural net initialization (I'm still a little afraid that the special cases I've seen in other truncated normal implementations are important, but I'm fairly confident that they're not needed for initializers).
• Adds a jax.nn namespace to centralize/official-ize the location of shared neural net-related functions.
• Puts implementations of standard initializers, activation functions, and a couple other things there. The goal is to have one place with reliable semantics and numerics; I'm not yet sure what kinds of tests to add to make sure that stays the case.
• Updates the log_softmax implementation to fix a numerical issue seen internally.
• Updates stax to use jax.nn, hopefully without breaking users.
• Adds docs for jax.nn.

Closes #1194 and #1195, re-closes #985.

[mattjj]

Nice! My main request is not to put code in __init__.py files and not to do from blah import *. A few other minor suggestions too.

[mattjj]

LGTM! Great idea to make this change. It's a big improvement and emblematic of a good way to think about jax-core's tools-for-nn-library-builders going forward.