Add Stateful Modules guide

[cgarciae]

What does this PR do?

WIP

Live Preview: https://flax--2689.org.readthedocs.build/en/2689/guides/stateful_modules.html

[codecov-commenter]

Codecov Report

Merging #2689 (ff95c1c) into main (db690a4) will increase coverage by 0.00%.
The diff coverage is n/a.

@@           Coverage Diff           @@
##             main    #2689   +/-   ##
=======================================
  Coverage   81.15%   81.15%           
=======================================
  Files          51       51           
  Lines        5492     5493    +1     
=======================================
+ Hits         4457     4458    +1     
  Misses       1035     1035           

Impacted Files	Coverage Δ
flax/errors.py	85.71% <0.00%> (+0.13%)	⬆️

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[marcvanzee]

Maybe we can add this to the "sharp bits" as well?

[8bitmp3]

TODO: @8bitmp3

[marcvanzee]

I would first explain that Flax Modules normally are stateless since Flax is using a functional API, and explain why this is the case. Then say that in some cases it can actually be useful to attach variables to the modules, and that the bind/unbind pattern allows this.

[cgarciae]

@jheek:

• Bound modules don't compose with any JAX API (cannot be tree_map-ed, vmap, etc).

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[cgarciae]

Currently bind fails for cases like this:

import flax.linen as nn
import jax
import jax.numpy as jnp

module = nn.Sequential([
    nn.Sequential([
        nn.Dense(8),
        nn.relu,
        nn.Dense(4),
    ]),
    nn.relu,
    nn.Dense(4),
])
x = jnp.ones((1, 5))
variables = module.init(jax.random.PRNGKey(0), x)
bound_module = module.bind(variables)

print(bound_module.layers[0].layers[0].scope) # None  <== !!!

@levskaya will try to come up with a solution to bind submodules, this guide will have to wait for that to be resolved. Previous attempt to solve this: #2028