Enhance JAX for the Impatient doc (Flax)

[8bitmp3]

A few small suggestions to improve the JAX for the Impatient doc:

• Provide a few minor fixes (grammar, syntax, spelling, capitalization (e.g. JVP) etc).
• At the beginning of the Gradients and autodiff section:
  • Spell out the JAX "docs" - the Autodiff Cookbook - for better UX.
  • Add "automatic differentiation system".

"For a full overview of JAX's automatic differentiation system, you can check the Autodiff Cookbook."

• Define VJPs and JVPs (Vector-Jacobian/Jacobian-Vector products) at the first mention.
  • Refer to them as reverse/forward-mode autodiff for less experienced users who haven't read the Autodiff Cookbook.
  • Provide specific links to VJP and JVP sections in the Cookbook, so the users don't have to scroll down to find them.
  • Capitalize "JVP", "VJP", "Jacobian-Vector" (as per the JAX glossary: https://jax.readthedocs.io/en/latest/glossary.html?highlight=JVP#term-JVP)

Even though, theoretically, a VJP (Vector-Jacobian product - reverse autodiff) and a JVP (Jacobian-Vector product - forward-mode autodiff) are similar—they compute a product of a Jacobian and a vector—they differ by the computational complexity of the operation. In short, when you have a large number of parameters (hence a wide matrix), a JVP is less efficient computationally than a VJP, and conversely JVP is more efficient when the Jacobian matrix is a tall matrix. You can read more in the JAX cookbook notebook mentioned above."

LMKWYT, thanks! @avital @marcvanzee @levskaya

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[8bitmp3]

The second warning in the pytest refers to the deprecated module, doesn't it?

WARNING: [autosummary] failed to import 'flax.linen.dot_product_attention_weights': no module named flax.linen.dot_product_attention_weights
11
WARNING: [autosummary] failed to import 'flax.optim.AdaFactor': no module named flax.optim.AdaFactor
12