Fixing small typo in docstring.

PiperOrigin-RevId: 466954884

Author: james-martens

kfac_jax/_src/optimizer.py

@@ -259,7 +259,7 @@ def __init__(
         specifying whether the batch is replicated over multiple devices and
         returns the batch size for a single device. (Default: ``None``)
       pmap_axis_name: String. The name of the pmap axis to use when
-        ``multi_device`` is set to True. (Default: ``curvature_axis``)
+        ``multi_device`` is set to True. (Default: ``kfac_axis``)
       forbid_setting_attributes_after_finalize: Boolean. By default after the
         object is finalized, you can not set any of its properties. This is done
         in order to protect the user from making changes to the object