FIX doctest format in TransformUpdateExtraArgsFn

PiperOrigin-RevId: 609652416

optax/_src/base.py

@@ -124,14 +124,14 @@ def __call__(
 
     For example, an update function that requires an additional loss parameter
     (which might be useful for implementing learning rate schedules that depend
-    on the current loss value) could be expressed as follows.
+    on the current loss value) could be expressed as follows:
 
     >>> def update(updates, state, params=None, *, loss, **extra_args):
-    >>>   del extra_args
-    >>>   # use loss value
+    ...   del extra_args
+    ...   # use loss value
 
-    Note that the loss value is keyword only, (it follows a `*` in the
-    signature of the function). This means users will get explicit errors if
+    Note that the loss value is keyword only, (it follows a ``*`` in the
+    signature of the function). This implies users will get explicit errors if
     they try to use this gradient transformation without providing the required
     argument.