Add Flax Dropout guide

[8bitmp3]

• Add a dedicated Dropout guide.
• Covers PRNG key usage, model creation (with nn.Module and nn.Dropout), initialization, the forward pass, and the training step (with TrainStep).
• Includes examples from flax/examples like a Transformer-based model trained on the WMT dataset (with dropout and attention dropout).

Shorten the Flax - The Sharp Bits by migrating some of the dropout instructions to the new guide. (To be done in a separate PR)

Live preview: https://flax--2675.org.readthedocs.build/en/2675/guides/dropout.html

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[marcvanzee]

Adding @cgarciae as a reviewer too since he's been thinking about this guide as well.

[8bitmp3]

Thanks @marcvanzee 👍

@cgarciae Your feedback would be super welcome 👍

[marcvanzee]

Great work! I think this will be an awesome tutorial, just needs a bit of work to get it really perfect.

[8bitmp3]

@marcvanzee @cgarciae Thanks. PTAL 👍

[cgarciae]

PTAL

[8bitmp3]

Thanks for the help @cgarciae 👍

@marcvanzee some of @cgarciae's review comments got stuck in "pending", they are out now (see above). We have gone over them and the new commits are on the way. Some changes will involve some changes to your code suggestions.

[8bitmp3]

@marcvanzee fyi working on adding the training step

[8bitmp3]

Reviewed the code and text with @marcvanzee. Committed and pushed the changes to the branch. cc @cgarciae. Amendments to the Sharp Bits will be in a separate PR.

[marcvanzee]

LGTM -- thanks a lot for this @8bitmp3, looks really awesome!! 🥳

(Please fix the doc build failures)

[codecov-commenter]

Codecov Report

Merging #2675 (c44cef0) into main (6d5bc2a) will not change coverage.
The diff coverage is n/a.

@@           Coverage Diff           @@
##             main    #2675   +/-   ##
=======================================
  Coverage   81.21%   81.21%           
=======================================
  Files          53       53           
  Lines        5605     5605           
=======================================
  Hits         4552     4552           
  Misses       1053     1053           

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

[8bitmp3]

All checks ✔️ @cgarciae PTAL. Thank you.

Live preview: https://flax--2675.org.readthedocs.build/en/2675/guides/dropout.html (build successful @marcvanzee).

[8bitmp3]

@cgarciae PTAL.

[cgarciae]

Thanks @8bitmp3! Looks great, lets merge this!