Small fixes to documentation.
Small fixes to documentation.

* Syntax error caused 'maybe update' to be missing (Fixes #185)
* Add some missing symbols from the public API to the docs
* Fix some small inconsistencies in docs (e.g. casing of titles)

PiperOrigin-RevId: 404493455

deepmind/optax,deepmind/optax