Our documentation for the modules provided by torch.nn is inconsistent and in some cases is missing information (e.g. shapes for forward() inputs). Since good documentation is crucial for the usability of these modules, this issue proposes to define a consistent template for module docs with all information necessary to use each effectively. What "template" means here could range from simply a cultural convention backed up by dev docs to technically enforced or generated docs (e.g. using jinja templates, CI checks, etc.).

Concretely, I propose that each module doc page should include documentation for:

• All public module constructor args, including expected shapes for all tensors
• All public forward() method inputs, including expected shapes for all tensors
• All forward() method outputs, including shapes for all tensors
• All learnable parameters associated with the module, including their shapes
• Any caveats, gotchas, limitations, etc.

To streamline the usage of these docs by users, the documentation sections should be consistently named and ordered & shape information should follow a consistent notation.

Open questions:

• How should the template be enforced? The loosest thing to do is enforce it by convention- i.e. document it with a full example in the developer docs & use this as a reference when adding new modules or modifying existing module docs. Alternatively, it may be possible to have a stricter enforcement by technical means - jinja templates, CI checks, etc.
• How well should this match the operator documentation? At a high level, modules can be thought of as just operators with added state. A lot of the same information provided in operator docs would make sense for module forward() docs. Examples: whether broadcasting is supported, supported dtypes, etc.
• Often, modules are simply stateful wrappers around functions in torch.nn.functional. What is the best way to identify this relationship in the docs? Perhaps there should be a section pointing to docs for the wrapped functional form, if there is one. See [docs] nn.modules pages should mention corresponding functional versions, e.g. for nn.Hardshrink #49992
• Anything missing above that should be in the docs?

cc @brianjo @mruberry @albanD @jbschlosser