Transformations
=====================

.. currentmodule:: optax

.. autosummary::
    :toctree: generated/

    adaptive_grad_clip
    AdaptiveGradClipState
    add_decayed_weights
    AddDecayedWeightsState
    add_noise
    AddNoiseState
    apply_every
    ApplyEvery
    bias_correction
    conditionally_mask
    conditionally_transform
    ConditionallyMaskState
    ConditionallyTransformState
    centralize
    clip
    clip_by_block_rms
    ClipState
    clip_by_global_norm
    ClipByGlobalNormState
    ema
    EmaState
    EmptyState
    global_norm
    GradientTransformation
    GradientTransformationExtraArgs
    identity
    init_empty_state
    keep_params_nonnegative
    measure_with_ema
    monitor
    MonitorState
    NonNegativeParamsState
    normalize_by_update_norm
    OptState
    Params
    per_example_global_norm_clip
    per_example_layer_norm_clip
    scale
    ScaleState
    scale_by_adadelta
    ScaleByAdaDeltaState
    scale_by_adan
    ScaleByAdanState
    scale_by_adam
    scale_by_adamax
    ScaleByAdamState
    scale_by_amsgrad
    ScaleByAmsgradState
    scale_by_backtracking_linesearch
    ScaleByBacktrackingLinesearchState
    scale_by_belief
    ScaleByBeliefState
    scale_by_factored_rms
    FactoredState
    scale_by_lbfgs
    ScaleByLBFGSState
    scale_by_learning_rate
    scale_by_lion
    ScaleByLionState
    scale_by_novograd
    ScaleByNovogradState
    scale_by_optimistic_gradient
    scale_by_param_block_norm
    scale_by_param_block_rms
    scale_by_polyak
    scale_by_radam
    scale_by_rms
    ScaleByRmsState
    scale_by_rprop
    ScaleByRpropState
    scale_by_rss
    ScaleByRssState
    scale_by_schedule
    ScaleByScheduleState
    scale_by_sign
    scale_by_sm3
    ScaleBySM3State
    scale_by_stddev
    ScaleByRStdDevState
    scale_by_trust_ratio
    ScaleByTrustRatioState
    scale_by_yogi
    scale_by_zoom_linesearch
    ScaleByZoomLinesearchState
    set_to_zero
    snapshot
    SnapshotState
    stateless
    stateless_with_tree_map
    trace
    TraceState
    TransformInitFn
    TransformUpdateFn
    TransformUpdateExtraArgsFn
    update_infinity_moment
    update_moment
    update_moment_per_elem_norm
    Updates
    with_extra_args_support
    zero_nans
    ZeroNansState
    ZoomLinesearchInfo
    masked
    freeze
    selective_transform