💡 README: Improve discoverability of Flax Basics 🔍

[8bitmp3]

Hi @avital @jheek @marcvanzee @levskaya @BertrandRdp @andsteing and the team/community👋

To help new users understand how certain key ML APIs/processes work with Flax during building and training ML models, I think we should update the main Flax README (under google/flax) by adding a few lines with code examples from the Flax basics notebook. Then, we could point to the docs with a short sentence like "To learn more about [optimizers/module basics/initialization/etc], check the Flax basics notebook".

For example, I think it's highly valuable to briefly mention in the README that:

• To perform gradient descent using Flax, you follow these steps: flax.optim -> wrapper with create() -> value_and_grad() etc.
• Model params arent' stored in a model. So, before you feed the params into an optimizer, you have to initialize them with JAX's random.PRNGKey() -> random.split() -> call model.init() etc.

The README is probably one of the first things that users come across when they hear about and check out Flax on GitHub. I've scratched my head while studying the code in linen_examples and the the Flax basics notebook really helped understand it better. Now, I think it'd be helpful if that doc is in the README — maybe somewhere at the top — as more users flock to google/flax.

LMKWYT 🙌

Happy holidays and stay safe 🍻

(from Don't Make Me Think - Steve Krug https://www.amazon.com/Dont-Make-Think-Revisited-Usability/dp/0321965515/)

[avital]

Hi @8bitmp3 -- and very very sorry for the late reply

I do agree that our README can be improved, and AFAIK nobody in particular is focusing on that at the moment. I'm happy to review any proposals you have for improvements to our README, our tutorials, documentation in general... It's best though if we can consider starting with smaller chunks of changes to get a sense of what you're proposing and debate them before any larger changes. I am really happy about what you did on #764 so I'm looking forward to seeing what other improvements we can make.

Let me know if there's any additional context or help that you could benefit from!