Minor fixes to docstrings.#376

copybara-service[bot] wants to merge 1 commit intomaingoogle-deepmind/kfac-jax:mainfrom

test_858685437google-deepmind/kfac-jax:test_858685437Copy head branch name to clipboard