Harmonize section order in docstrings

[mblondel]

My preference goes to: Args, Returns, Example, References

[fabianp]

it would also be important to add whatever is decided on the development docs (https://optax.readthedocs.io/en/latest/development.html)