Make documentation examples runnable without changes

UPDATE: Updated title based on discussion to make examples runnable without changes as opposed to just including imports.

Many of the self contained examples in the documentation do not include imports, which can make getting started tough for new users. For example, many examples use nn.Module, and a new user may try from flax import nn instead of import flax.linen as nn when trying to run the example on their own. I think all standalone snippets should include all necessary imports, so one can simply paste the example and run.

Here’s an initial list of examples I think could use imports:

• All the readme examples (these are quite inconsistent, e.g. this line uses Dense directly instead of nn.Dense)
• Managing Parameters and State
• Extracting Intermediates
• Model Surgery
• Module.bind
• Module.sow

I noticed Flax uses doctest to ensure examples in the documentation run, is there any way to mark/enforce that certain examples are self-contained (i.e. have the required imports)?

cc @8bitmp3

[avital]

As for from flax import nn - @avital why don’t we currently have a deprecation warning on import?

It’s because we re-export flax.nn in flax/__init__.py so that import flax flax.nn.<...> works. But that seems like a bad answer. We should get rid of that, but honestly it’s probably time to remove flax.nn.

[andsteing]

Actually, these examples are run with the Sphinx doctest extension:

https://github.com/google/flax/blob/491ce18759622506588784b4fca0e4bf05f8c8cd/docs/conf.py#L54

Maybe we should transform the testsetup blocks into testcode blocks so the imports are shown in the docs as well?

@n2cholas would you really want every snippet to contain all the imports or would it be enough to have a list of imports at the top of the file? The first might be a bit too verbose and it should be easy enough to copy the imports from the same page when copying the snippet…

As for from flax import nn - @avital why don’t we currently have a deprecation warning on import? Because that indeed seems like a thing a beginner would do, and it might be rather confusing to have to discover that “Linen is the new nn”.