Neural network parameters should be public and better documented

[nathanielsimard]

Feature description

Neural network modules in burn-core/src/nn/* should have their parameters made public and documented.

Feature motivation

Having access to module parameters improves the debugging process by making it trivial to print tensors and gradients. It also improves the flexibility of the framework, since you can use the parameters the way you want/need.

Suggest a Solution

We should think a bit about where the documentation is going. Here's a list of items that every module should have.

• [ ] Parameter initialization methods
• [ ] Parameter shapes
• [ ] Config/Hyper-parameter descriptions
• [ ] Math formula of the forward pass
• [ ] Global description

Since Burn has a separation between Config and the module struct, maybe we should put the description on top of the most relevant field. The initialization scheme can be put on top of initializer, usually in the Config struct, the shapes can be put on top of each parameter field in the module struct (e.g., weight: Param<Tensor<B, 2>>), the math formulas can be put on top of the forward pass and/or the global description, etc. I think it's important to link each part in the global description in a structured manner to quickly have the information we want.

[camdenmcgath]

@nathanielsimard Is this issue still needing to be fixed? I see that it is open, so I'll get working on the documentation so long as no one calls me off...

[nathanielsimard]

Yes this is still open, the goal is to offer great documentation for nn modules, better than what we have right now!

[onurcanbektas]

@camdenmcgath Are you still working on this?

[antimora]

@onurcanbektas, it is available if you like to jump in. Some of NN module were corrected but it was by need.

[laggui]

Completed in #1746