Fix docs errors

PiperOrigin-RevId: 626833460
google-deepmind · Apr 23, 2024 · 9c7a290 · 9c7a290
1 parent 748ce7f
commit 9c7a290
Show file tree

Hide file tree

Showing 20 changed files with 131 additions and 200 deletions.
diff --git a/docs/api/combining_optimizers.rst b/docs/api/combining_optimizers.rst
@@ -15,4 +15,3 @@ Multi-transform
 ~~~~~~~~~~~~~~~
 .. autofunction:: multi_transform
 .. autoclass::  MultiTransformState
-   :members:
diff --git a/docs/api/contrib.rst b/docs/api/contrib.rst
@@ -27,50 +27,34 @@ Complex-valued Optimization
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~
 .. autofunction:: split_real_and_imaginary
 .. autoclass:: SplitRealAndImaginaryState
-    :members:
 
 Continuous coin betting
 ~~~~~~~~~~~~~~~~~~~~~~~
 .. autofunction:: cocob
 .. autoclass:: COCOBState
-    :members:
 
 D-adaptation
 ~~~~~~~~~~~~
 .. autofunction:: dadapt_adamw
 .. autoclass:: DAdaptAdamWState
-    :members:
-
-Privacy-Sensitive Optax Methods
--------------------------------
-
-.. autosummary::
-    DifferentiallyPrivateAggregateState
-    differentially_private_aggregate
-
 
 Differentially Private Aggregate
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 .. autofunction:: differentially_private_aggregate
 .. autoclass:: DifferentiallyPrivateAggregateState
-   :members:
 .. autofunction:: dpsgd
 
-
 Mechanize
 ~~~~~~~~~
 .. autofunction:: mechanize
 .. autoclass:: MechanicState
-    :members:
 
 Prodigy
 ~~~~~~~
 .. autofunction:: prodigy
 .. autoclass:: ProdigyState
-    :members:
 
 Sharpness aware minimization
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 .. autofunction:: sam
 .. autoclass:: SAMState
-    :members:
diff --git a/docs/api/control_variates.rst b/docs/api/control_variates.rst
@@ -9,13 +9,13 @@ Control Variates
     moving_avg_baseline
 
 Control delta method
-~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~
 .. autofunction:: control_delta_method
 
 Control variates Jacobians
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~~~~~
 .. autofunction:: control_variates_jacobians
 
 Moving average baseline
-~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~~
 .. autofunction:: moving_avg_baseline
diff --git a/docs/api/losses.rst b/docs/api/losses.rst
@@ -54,7 +54,7 @@ Kullback-Leibler divergence
 .. autofunction:: kl_divergence
 
 L2 Squared loss
-~~~~~~~~~~~~
+~~~~~~~~~~~~~~~
 .. autofunction:: squared_error
 .. autofunction:: l2_loss
 

diff --git a/docs/api/optimizer_schedules.rst b/docs/api/optimizer_schedules.rst
@@ -23,7 +23,6 @@ Optimizer Schedules
 
 
 .. autoclass:: Schedule
-   :members:
 
 Constant schedule
 ~~~~~~~~~~~~~~~~~
@@ -46,7 +45,6 @@ Inject hyperparameters
 ~~~~~~~~~~~~~~~~~~~~~~
 .. autofunction:: inject_hyperparams
 .. autoclass:: InjectHyperparamsState
-   :members:
 
 Linear schedules
 ~~~~~~~~~~~~~~~~

diff --git a/docs/api/optimizer_wrappers.rst b/docs/api/optimizer_wrappers.rst
@@ -22,42 +22,34 @@ Optimizer Wrappers
 
 
 Apply if finite
-~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~
 .. autofunction::  apply_if_finite
 .. autoclass::  ApplyIfFiniteState
-   :members:
 
 Flatten
 ~~~~~~~~
 .. autofunction:: flatten
 
 Lookahead
-~~~~~~~~~~~~~~~~~
+~~~~~~~~~
 .. autofunction::  lookahead
 .. autoclass::  LookaheadParams
-   :members:
 .. autoclass::  LookaheadState
-   :members:
 
 Masked update
-~~~~~~~~~~~~~~
+~~~~~~~~~~~~~
 .. autofunction::  masked
 .. autoclass::  MaskedState
-   :members:
 
 Maybe update
-~~~~~~~~~~~~~~
+~~~~~~~~~~~~
 .. autofunction:: maybe_update
 .. autoclass:: MaybeUpdateState
-   :members:
 
 Multi-step update
-~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~
 .. autoclass:: MultiSteps
-   :members:
 .. autoclass:: MultiStepsState
-   :members:
 .. autoclass:: ShouldSkipUpdateFunction
-   :members:
 .. autofunction:: skip_large_updates
 .. autofunction:: skip_not_finite
diff --git a/docs/api/transformations.rst b/docs/api/transformations.rst
@@ -96,122 +96,96 @@ Types
 ~~~~~
 
 .. autoclass:: GradientTransformation
-    :members:
 
 .. autoclass:: GradientTransformationExtraArgs
-    :members:
 
 .. autoclass:: TransformInitFn
-    :members:
 
 .. autoclass:: TransformUpdateFn
-    :members:
 
 .. autoclass:: OptState
-    :members:
 
 .. autoclass:: Params
-    :members:
 
 .. autoclass:: Updates
-    :members:
 
 
 Transformations and states
 ~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 .. autofunction:: adaptive_grad_clip
 .. autoclass:: AdaptiveGradClipState
-  :members:
 
 .. autofunction:: add_decayed_weights
 .. autoclass:: AddDecayedWeightsState
-    :members:
 
 .. autofunction:: add_noise
 .. autoclass:: AddNoiseState
-    :members:
 
 .. autofunction:: apply_every
 .. autoclass:: ApplyEvery
-    :members:
 
 .. autofunction:: bias_correction
 
 .. autofunction:: centralize
 
 .. autofunction:: conditionally_mask
 .. autoclass:: ConditionallyMaskState
-    :members:
 
 .. autofunction:: conditionally_transform
 .. autoclass:: ConditionallyTransformState
-    :members:
 
 .. autofunction:: clip
 .. autofunction:: clip_by_block_rms
 .. autoclass:: ClipState
-    :members:
 
 .. autofunction:: clip_by_global_norm
 .. autoclass:: ClipByGlobalNormState
-    :members:
 
 .. autofunction:: ema
 .. autoclass:: EmaState
-    :members:
 
 .. autoclass:: EmptyState
-    :members:
 
 .. autofunction:: global_norm
 
 .. autofunction:: identity
 
 .. autofunction:: keep_params_nonnegative
 .. autoclass:: NonNegativeParamsState
-    :members:
 
 .. autofunction:: per_example_global_norm_clip
 .. autofunction:: per_example_layer_norm_clip
 
 .. autofunction:: scale
 .. autoclass:: ScaleState
-    :members:
 
 .. autofunction:: scale_by_adadelta
 .. autoclass:: ScaleByAdaDeltaState
-    :members:
 
 .. autofunction:: scale_by_adam
 .. autofunction:: scale_by_adamax
 .. autoclass:: ScaleByAdamState
-    :members:
 
 .. autofunction:: scale_by_amsgrad
 .. autoclass:: ScaleByAmsgradState
-    :members:
 
 .. autofunction:: scale_by_backtracking_linesearch
 .. autoclass:: ScaleByBacktrackingLinesearchState
 
 .. autofunction:: scale_by_belief
 .. autoclass:: ScaleByBeliefState
-    :members:
 
 .. autofunction:: scale_by_factored_rms
 .. autoclass:: FactoredState
-    :members:
 
 .. autofunction:: scale_by_learning_rate
 
 .. autofunction:: scale_by_lion
 .. autoclass:: ScaleByLionState
-    :members:
 
 .. autofunction:: scale_by_novograd
 .. autoclass:: ScaleByNovogradState
-    :members:
 
 .. autofunction:: scale_by_optimistic_gradient
 
@@ -225,31 +199,24 @@ Transformations and states
 
 .. autofunction:: scale_by_rms
 .. autoclass:: ScaleByRmsState
-    :members:
 
 .. autofunction:: scale_by_rprop
 .. autoclass:: ScaleByRpropState
-    :members:
 
 .. autofunction:: scale_by_rss
 .. autoclass:: ScaleByRssState
-    :members:
 
 .. autofunction:: scale_by_schedule
 .. autoclass:: ScaleByScheduleState
-    :members:
 
 .. autofunction:: scale_by_sm3
 .. autoclass:: ScaleBySM3State
-    :members:
 
 .. autofunction:: scale_by_stddev
 .. autoclass:: ScaleByRStdDevState
-    :members:
 
 .. autofunction:: scale_by_trust_ratio
 .. autoclass:: ScaleByTrustRatioState
-    :members:
 
 .. autofunction:: scale_by_yogi
 
@@ -260,7 +227,6 @@ Transformations and states
 
 .. autofunction:: trace
 .. autoclass:: TraceState
-    :members:
 
 .. autofunction:: update_infinity_moment
 .. autofunction:: update_moment
@@ -270,4 +236,3 @@ Transformations and states
 
 .. autofunction:: zero_nans
 .. autoclass:: ZeroNansState
-    :members: