jax

e9909ce0 - Copybara import of the project:

Copybara import of the project: -- a396cfbbd414f6f21f0c7e8a68e6e89d202c0e84 by Peter Hawkins <phawkins@google.com>: Move jax.nn implementation into jax._src.nn. PiperOrigin-RevId: 337671917