Format code base.

[copybara-service]

Format code base.

Fixes #1112, fixes #1111

• Formatting all files with internal formatter
• Use correct rst format for references all over the files
• Refer to :class:optax.GradientTransformation all over the code rather than GradientTransformation
• Use note, warnings in rst format all over the files (rather than some NOTE:... or WARNING:... in the files)
• Apply following template for all files (and added this template to the docs)
  One-line doc
  Description
  Args
  Returns
  Examples
  References
  See also, note, warning, etc...