Add missing members to the doc

  * Some classes and functions were mentioned in the "autosummary" section but didn't have an associated autoclass/autofunction

  * refactor utilities.rst so it conforms to the same structure as other pages, with an initial summary of members followed by a more detailed description

PiperOrigin-RevId: 601363946

docs/api/optimizer_wrappers.rst

@@ -57,3 +57,7 @@ Multi-step update
    :members:
 .. autoclass:: MultiStepsState
    :members:
+.. autoclass:: ShouldSkipUpdateFunction
+   :members:
+.. autofunction:: skip_large_updates
+.. autofunction:: skip_not_finite

docs/api/transformations.rst

@@ -79,12 +79,15 @@ Transformations
     with_extra_args_support
 
 
-Optax types
-~~~~~~~~~~~~~~
+Types
+~~~~~
 
 .. autoclass:: GradientTransformation
     :members:
 
+.. autoclass:: GradientTransformationExtraArgs
+    :members:
+
 .. autoclass:: TransformInitFn
     :members:
 
@@ -101,8 +104,8 @@ Optax types
     :members:
 
 
-Optax transformations and states
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Transformations and states
+~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 .. autofunction:: adaptive_grad_clip
 .. autoclass:: AdaptiveGradClipState
@@ -120,6 +123,8 @@ Optax transformations and states
 .. autoclass:: ApplyEvery
     :members:
 
+.. autofunction:: bias_correction
+
 .. autofunction:: centralize
 
 .. autofunction:: clip
@@ -219,6 +224,9 @@ Optax transformations and states
 .. autoclass:: TraceState
     :members:
 
+.. autofunction:: update_infinity_moment
+.. autofunction:: update_moment
+.. autofunction:: update_moment
 .. autofunction:: zero_nans
 .. autoclass:: ZeroNansState
     :members:

docs/api/utilities.rst

@@ -1,13 +1,34 @@
 Utilities
 =========
 
-General
--------
-
 .. currentmodule:: optax
 
 .. autosummary::
     scale_gradient
+    safe_int32_increment
+    safe_norm
+    safe_root_mean_squares
+    DifferentiallyPrivateAggregateState
+    differentially_private_aggregate
+    fisher_diag
+    hessian_diag
+    hvp
+    tree_add
+    tree_add_scalar_mul
+    tree_div
+    tree_vdot
+    tree_l2_norm
+    tree_map_params
+    tree_mul
+    tree_ones_like
+    tree_scalar_mul
+    tree_sub
+    tree_sum
+    tree_zeros_like
+
+
+General
+-------
 
 Scale gradient
 ~~~~~~~~~~~~~~
@@ -17,13 +38,6 @@ Scale gradient
 Numerical Stability
 -------------------
 
-.. currentmodule:: optax
-
-.. autosummary::
-    safe_int32_increment
-    safe_norm
-    safe_root_mean_squares
-
 Safe int32 increment
 ~~~~~~~~~~~~~~~~~~~~
 .. autofunction:: safe_int32_increment
@@ -37,16 +51,6 @@ Safe root mean squares
 .. autofunction:: safe_root_mean_squares
 
 
-Privacy-Sensitive Optax Methods
--------------------------------
-
-.. currentmodule:: optax
-
-.. autosummary::
-    DifferentiallyPrivateAggregateState
-    differentially_private_aggregate
-
-
 Differentially Private Aggregate
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 .. autofunction:: differentially_private_aggregate
@@ -77,24 +81,8 @@ Hessian vector product
 .. autofunction:: hvp
 
 
-Tree
-----
-
-.. currentmodule:: optax.tree_utils
-
-.. autosummary::
-    tree_add
-    tree_add_scalar_mul
-    tree_div
-    tree_vdot
-    tree_l2_norm
-    tree_map_params
-    tree_mul
-    tree_ones_like
-    tree_scalar_mul
-    tree_sub
-    tree_sum
-    tree_zeros_like
+Tree manipulation
+-----------------
 
 Tree add
 ~~~~~~~~