doc suggestions

[benjaminpope]

Hi Louis,

Great work.

Suggestion about the docs: the huge long notebook could be split up. I think a set of very brief notebooks showing

• get, set, add, pmap
• optax
• numpyro APIs would be more accessible for a general reader.

The text in the first paragraph at the top could be edited for tone and length and go onto the landing page rather than one of the notebooks.

Cheers,

Ben

[LouisDesdoigts]

Done with v0.3.0