"""
@generated by mypy-protobuf.  Do not edit manually!
isort:skip_file
"""

import builtins
import sys
import typing

import google.protobuf.descriptor
import google.protobuf.internal.enum_type_wrapper
import google.protobuf.message
import google.protobuf.wrappers_pb2
import tensorflow.compiler.xla.service.hlo_pb2

if sys.version_info >= (3, 10):
    import typing as typing_extensions
else:
    import typing_extensions

DESCRIPTOR: google.protobuf.descriptor.FileDescriptor

@typing.final
class ClippingLimits(google.protobuf.message.Message):
    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    LOWER_FIELD_NUMBER: builtins.int
    UPPER_FIELD_NUMBER: builtins.int
    @property
    def lower(self) -> google.protobuf.wrappers_pb2.FloatValue:
        """-inf if not set"""

    @property
    def upper(self) -> google.protobuf.wrappers_pb2.FloatValue:
        """+inf if not set"""

    def __init__(
        self,
        *,
        lower: google.protobuf.wrappers_pb2.FloatValue | None = ...,
        upper: google.protobuf.wrappers_pb2.FloatValue | None = ...,
    ) -> None: ...
    def HasField(self, field_name: typing.Literal["lower", b"lower", "upper", b"upper"]) -> builtins.bool: ...
    def ClearField(self, field_name: typing.Literal["lower", b"lower", "upper", b"upper"]) -> None: ...

global___ClippingLimits = ClippingLimits

@typing.final
class SimulatedQuantization(google.protobuf.message.Message):
    """Configuration for simulated quantization; simulated quantization is used to
    reduce training/serving skew when the serving variables are quantized. The
    same quantization operations are executed during training to minimize
    differences with serving.

    Simulated quantization inserts the following operations on the forward pass
    after gathering the embedding vector from HBM. The backward pass operations
    are unchanged.

    clipped_val = clip(input, clipping_limits)
    quantum = clipping_limits.range() / (num_buckets - 1)
    quantized_val = floor((clipped_val - clipping_limits.lower()) / quantum + .5)
    return quantized_val * quantum + clipping_limits.lower().
    """

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    ENABLED_FIELD_NUMBER: builtins.int
    CLIPPING_LIMITS_FIELD_NUMBER: builtins.int
    NUM_BUCKETS_FIELD_NUMBER: builtins.int
    enabled: builtins.bool
    """Whether simulated quantization is enabled."""
    num_buckets: builtins.int
    """Number of possible quantized values."""
    @property
    def clipping_limits(self) -> global___ClippingLimits:
        """Minimum and maximum values of the range used for quantization."""

    def __init__(
        self,
        *,
        enabled: builtins.bool | None = ...,
        clipping_limits: global___ClippingLimits | None = ...,
        num_buckets: builtins.int | None = ...,
    ) -> None: ...
    def HasField(self, field_name: typing.Literal["clipping_limits", b"clipping_limits"]) -> builtins.bool: ...
    def ClearField(
        self,
        field_name: typing.Literal["clipping_limits", b"clipping_limits", "enabled", b"enabled", "num_buckets", b"num_buckets"],
    ) -> None: ...

global___SimulatedQuantization = SimulatedQuantization

@typing.final
class DynamicLearningRate(google.protobuf.message.Message):
    """Dynamic learning rate specification in the TPUEmbeddingConfiguration. The
    actual learning rates are provided as a scalar input list to the
    SendTPUEmbeddingGradients Op indexed by their tag specified through the
    following proto.
    """

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    TAG_FIELD_NUMBER: builtins.int
    tag: builtins.int
    """For tables where learning rates are dynamically computed and communicated
    to the TPU embedding program, a tag must be specified for the learning
    rate.

    The tag must be a non-negative  integer. The total number of unique tags
    must be less than or equal to the number of tables in the TPU embedding
    configuration (a table does not specify any tag if it uses a constant
    learning rate, and specifies exactly one tag if it uses dynamic learning
    rates).

    All tags in the range [0, number_of_unique_tags) must be present in the TPU
    embedding configuration, i.e. a tag cannot be skipped if a different tag
    numerically greater than it is used in the configuration.

    If multiple tables specify the same tag, they *MUST* have
    the same dynamic learning rate, for example, their dynamic learning rate
    could be computed by the same TensorFlow sub-graph. The partitioning of the
    embedding layer would be more optimal if the number_of_unique_tags is as
    *LOW* as possible, i.e., if many tables share the same tag.

    The learning_rate input of the SendTPUEmbeddingGradients op is used to
    communicate dynamic learning rates to the TPU embedding program.
    The learning_rate input is a list of scalars where the size of the list is
    equal to the number of unique tags. The learning rate associated with a
    particular tag is specified by populating its corresponding index in the
    list of learning_rate scalars.
    """
    def __init__(self, *, tag: builtins.int | None = ...) -> None: ...
    def ClearField(self, field_name: typing.Literal["tag", b"tag"]) -> None: ...

global___DynamicLearningRate = DynamicLearningRate

@typing.final
class LearningRate(google.protobuf.message.Message):
    """Source of learning rate to use."""

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    CONSTANT_FIELD_NUMBER: builtins.int
    DYNAMIC_FIELD_NUMBER: builtins.int
    constant: builtins.float
    @property
    def dynamic(self) -> global___DynamicLearningRate: ...
    def __init__(self, *, constant: builtins.float | None = ..., dynamic: global___DynamicLearningRate | None = ...) -> None: ...
    def HasField(
        self, field_name: typing.Literal["constant", b"constant", "dynamic", b"dynamic", "learning_rate", b"learning_rate"]
    ) -> builtins.bool: ...
    def ClearField(
        self, field_name: typing.Literal["constant", b"constant", "dynamic", b"dynamic", "learning_rate", b"learning_rate"]
    ) -> None: ...
    def WhichOneof(
        self, oneof_group: typing.Literal["learning_rate", b"learning_rate"]
    ) -> typing.Literal["constant", "dynamic"] | None: ...

global___LearningRate = LearningRate

@typing.final
class AdagradParameters(google.protobuf.message.Message):
    """Each optimizer's parameter proto has a link to its documentation and CPU
    implementation (if available) for user reference.

    https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adagrad
    https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L1634
    """

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    def __init__(self) -> None: ...

global___AdagradParameters = AdagradParameters

@typing.final
class AdagradMomentumParameters(google.protobuf.message.Message):
    """This optimizer combines the Adagrad and Momentum update rules.
    accum(new) = beta2 == 1.0 ?
                 accum(old) + grad^2 :
                 beta2 * accum(old) + (1 - beta2) * grad^2
    accum_with_exponent = (accum(new) + epsilon)^(-1.0 / exponent)
    mom_accum(new) = momentum * mom_accum(old) + accum_with_exponent
    update = use_nesterov ?
             momentum * mom_accum(new) + accum_with_exponent :
             mom_accum(new)
    var(new) = var(old) - lr * grad * update
    Algorithm described in https://arxiv.org/abs/2002.11803.
    """

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    MOMENTUM_FIELD_NUMBER: builtins.int
    USE_NESTEROV_FIELD_NUMBER: builtins.int
    EXPONENT_FIELD_NUMBER: builtins.int
    BETA2_FIELD_NUMBER: builtins.int
    EPSILON_FIELD_NUMBER: builtins.int
    momentum: builtins.float
    """Moving average parameter for the momentum accumulator."""
    use_nesterov: builtins.bool
    """Whether to use the Nesterov variant of momentum."""
    exponent: builtins.float
    """Exponent for the gradient^2 accumulator."""
    beta2: builtins.float
    """Moving average parameter for the gradient^2 accumulator."""
    epsilon: builtins.float
    """Offset added to the Adagrad accumulator."""
    def __init__(
        self,
        *,
        momentum: builtins.float | None = ...,
        use_nesterov: builtins.bool | None = ...,
        exponent: builtins.float | None = ...,
        beta2: builtins.float | None = ...,
        epsilon: builtins.float | None = ...,
    ) -> None: ...
    def ClearField(
        self,
        field_name: typing.Literal[
            "beta2",
            b"beta2",
            "epsilon",
            b"epsilon",
            "exponent",
            b"exponent",
            "momentum",
            b"momentum",
            "use_nesterov",
            b"use_nesterov",
        ],
    ) -> None: ...

global___AdagradMomentumParameters = AdagradMomentumParameters

@typing.final
class BoundedAdagradParameters(google.protobuf.message.Message):
    """Algorithm in http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf."""

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    UPDATE_ACCUMULATOR_FIRST_FIELD_NUMBER: builtins.int
    MAX_VAR_UPDATE_FIELD_NUMBER: builtins.int
    MAX_ACCUMULATOR_FIELD_NUMBER: builtins.int
    update_accumulator_first: builtins.bool
    """Whether to use the updated or the old value of the accumulator when
    computing the effective learning rate. When update_accumulator_first is set
    to True, the updated value of the accumulator is used.
    """
    max_var_update: builtins.float
    """The max_var_update value to use. Set value to 0 (default) to disable using
    max_var_update to clip the gradient.
    """
    max_accumulator: builtins.float
    """The maximum value of the accumulator. Set max_accumulator to 0 (default)
    to disable using max_accumulator to clip the accumulator.
    """
    def __init__(
        self,
        *,
        update_accumulator_first: builtins.bool | None = ...,
        max_var_update: builtins.float | None = ...,
        max_accumulator: builtins.float | None = ...,
    ) -> None: ...
    def ClearField(
        self,
        field_name: typing.Literal[
            "max_accumulator",
            b"max_accumulator",
            "max_var_update",
            b"max_var_update",
            "update_accumulator_first",
            b"update_accumulator_first",
        ],
    ) -> None: ...

global___BoundedAdagradParameters = BoundedAdagradParameters

@typing.final
class StochasticGradientDescentParameters(google.protobuf.message.Message):
    """https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/SGD
    https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L629
    """

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    def __init__(self) -> None: ...

global___StochasticGradientDescentParameters = StochasticGradientDescentParameters

@typing.final
class FtrlParameters(google.protobuf.message.Message):
    """https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Ftrl
    https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/41159.pdf
    https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L2646

    The hyperparameters for FTRL are the same as for the Keras implementation,
    with some additions. The "beta" parameter matches the behavior described in
    the second link above; "beta" / (2 * learning rate) should be added to "l2"
    to get equivalent behavior in the other TensorFlow implementations of this
    optimizer. When the multiply_linear_by_lr field is set to true, a modified
    formula is used for FTRL that treats the "linear" accumulator as being
    pre-multiplied by the learning rate (i.e., the accumulator named "linear"
    actually stores "linear * learning_rate"). Other than checkpoint
    compatibility, this is mathematically equivalent for a static learning rate;
    for a dynamic learning rate, it is nearly the same as long as the learning
    rate does not change quickly. The benefit of setting multiply_linear_by_lr to
    true is that the modified formula handles zero and near-zero learning rates
    without producing NaNs, improving flexibility for learning rate ramp-up.
    """

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    L1_FIELD_NUMBER: builtins.int
    L2_FIELD_NUMBER: builtins.int
    LR_POWER_FIELD_NUMBER: builtins.int
    BETA_FIELD_NUMBER: builtins.int
    MULTIPLY_LINEAR_BY_LR_FIELD_NUMBER: builtins.int
    ALLOW_ZERO_ACCUMULATOR_FIELD_NUMBER: builtins.int
    l1: builtins.float
    l2: builtins.float
    lr_power: builtins.float
    beta: builtins.float
    multiply_linear_by_lr: builtins.bool
    allow_zero_accumulator: builtins.bool
    """Previously, allow_zero_accumulator parameter changed some internal formulas
    to allow zero and near-zero accumulator values at the cost of some
    performance. The current implementation ignores this parameter; zero or
    near-zero accumulator values are now always supported.
    """
    def __init__(
        self,
        *,
        l1: builtins.float | None = ...,
        l2: builtins.float | None = ...,
        lr_power: builtins.float | None = ...,
        beta: builtins.float | None = ...,
        multiply_linear_by_lr: builtins.bool | None = ...,
        allow_zero_accumulator: builtins.bool | None = ...,
    ) -> None: ...
    def ClearField(
        self,
        field_name: typing.Literal[
            "allow_zero_accumulator",
            b"allow_zero_accumulator",
            "beta",
            b"beta",
            "l1",
            b"l1",
            "l2",
            b"l2",
            "lr_power",
            b"lr_power",
            "multiply_linear_by_lr",
            b"multiply_linear_by_lr",
        ],
    ) -> None: ...

global___FtrlParameters = FtrlParameters

@typing.final
class AdamParameters(google.protobuf.message.Message):
    """The Adam optimizer does not implement hyper-parameter update due to hardware
    limitations; use the dynamic learning rate feature instead, setting the
    learning rate to: user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
    Here, t is the current timestep.

    https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adam
    https://github.com/tensorflow/tensorflow/blob/ab51450c817674c8ff08a7ae4f8ac50cdc4bed8b/tensorflow/python/training/adam.py#L32

    Note that the code by default implements the lazy version of Adam
    (https://www.tensorflow.org/api_docs/python/tf/contrib/opt/LazyAdamOptimizer)
    unless the use_non_lazy_adam parameter is set, in which case it implements
    the normal version of Adam that updates all parameters in the embedding
    table, even for entries that are not used in the current minibatch
    (https://www.tensorflow.org/api_docs/python/tf/contrib/opt/AdamOptimizer). If
    use_non_lazy_adam is enabled, gradient accumulation is also required to be
    enabled in order to get correct results; a warning will be printed otherwise
    (which may change to an error in the future). If use_sum_inside_sqrt is set,
    the Adam variable update formula will be changed from m / (sqrt(v) + epsilon)
    to m / sqrt(v + epsilon**2); this option improves the performance of TPU
    training and is not expected to harm model quality.
    """

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    BETA1_FIELD_NUMBER: builtins.int
    BETA2_FIELD_NUMBER: builtins.int
    EPSILON_FIELD_NUMBER: builtins.int
    USE_NON_LAZY_ADAM_FIELD_NUMBER: builtins.int
    USE_SUM_INSIDE_SQRT_FIELD_NUMBER: builtins.int
    beta1: builtins.float
    beta2: builtins.float
    epsilon: builtins.float
    use_non_lazy_adam: builtins.bool
    use_sum_inside_sqrt: builtins.bool
    def __init__(
        self,
        *,
        beta1: builtins.float | None = ...,
        beta2: builtins.float | None = ...,
        epsilon: builtins.float | None = ...,
        use_non_lazy_adam: builtins.bool | None = ...,
        use_sum_inside_sqrt: builtins.bool | None = ...,
    ) -> None: ...
    def ClearField(
        self,
        field_name: typing.Literal[
            "beta1",
            b"beta1",
            "beta2",
            b"beta2",
            "epsilon",
            b"epsilon",
            "use_non_lazy_adam",
            b"use_non_lazy_adam",
            "use_sum_inside_sqrt",
            b"use_sum_inside_sqrt",
        ],
    ) -> None: ...

global___AdamParameters = AdamParameters

@typing.final
class MomentumParameters(google.protobuf.message.Message):
    """https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/SGD
    https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L3068
    """

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    MOMENTUM_FIELD_NUMBER: builtins.int
    USE_NESTEROV_FIELD_NUMBER: builtins.int
    momentum: builtins.float
    use_nesterov: builtins.bool
    def __init__(self, *, momentum: builtins.float | None = ..., use_nesterov: builtins.bool | None = ...) -> None: ...
    def ClearField(self, field_name: typing.Literal["momentum", b"momentum", "use_nesterov", b"use_nesterov"]) -> None: ...

global___MomentumParameters = MomentumParameters

@typing.final
class LionParameters(google.protobuf.message.Message):
    """https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Lion
    momenta(new) = beta2 * momenta(old) + (1 - beta2) * grad
    momenta_t = beta1 * momenta(old) + (1 - beta1) * grad
    var(new) = var(old) - lr * sign(momenta_t)
    Algorithm described in https://arxiv.org/abs/2302.06675.
    """

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    BETA1_FIELD_NUMBER: builtins.int
    BETA2_FIELD_NUMBER: builtins.int
    USE_NON_LAZY_LION_FIELD_NUMBER: builtins.int
    beta1: builtins.float
    beta2: builtins.float
    use_non_lazy_lion: builtins.bool
    def __init__(
        self,
        *,
        beta1: builtins.float | None = ...,
        beta2: builtins.float | None = ...,
        use_non_lazy_lion: builtins.bool | None = ...,
    ) -> None: ...
    def ClearField(
        self, field_name: typing.Literal["beta1", b"beta1", "beta2", b"beta2", "use_non_lazy_lion", b"use_non_lazy_lion"]
    ) -> None: ...

global___LionParameters = LionParameters

@typing.final
class RmsPropParameters(google.protobuf.message.Message):
    """https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/RMSprop
    https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L4229
    """

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    RHO_FIELD_NUMBER: builtins.int
    MOMENTUM_FIELD_NUMBER: builtins.int
    EPSILON_FIELD_NUMBER: builtins.int
    rho: builtins.float
    momentum: builtins.float
    epsilon: builtins.float
    def __init__(
        self, *, rho: builtins.float | None = ..., momentum: builtins.float | None = ..., epsilon: builtins.float | None = ...
    ) -> None: ...
    def ClearField(self, field_name: typing.Literal["epsilon", b"epsilon", "momentum", b"momentum", "rho", b"rho"]) -> None: ...

global___RmsPropParameters = RmsPropParameters

@typing.final
class CenteredRmsPropParameters(google.protobuf.message.Message):
    """https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/RMSprop
    https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L4358
    """

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    RHO_FIELD_NUMBER: builtins.int
    MOMENTUM_FIELD_NUMBER: builtins.int
    EPSILON_FIELD_NUMBER: builtins.int
    rho: builtins.float
    momentum: builtins.float
    epsilon: builtins.float
    def __init__(
        self, *, rho: builtins.float | None = ..., momentum: builtins.float | None = ..., epsilon: builtins.float | None = ...
    ) -> None: ...
    def ClearField(self, field_name: typing.Literal["epsilon", b"epsilon", "momentum", b"momentum", "rho", b"rho"]) -> None: ...

global___CenteredRmsPropParameters = CenteredRmsPropParameters

@typing.final
class MdlAdagradLightParameters(google.protobuf.message.Message):
    """Variant of algorithm in http://proceedings.mlr.press/v44/shamir15.pdf"""

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    L2_FIELD_NUMBER: builtins.int
    LR_POWER_FIELD_NUMBER: builtins.int
    MIN_SERVABLE_MDL_BENEFIT_FIELD_NUMBER: builtins.int
    MDL_MIX_IN_MARGIN_FIELD_NUMBER: builtins.int
    MDL_BENEFIT_RAMPUP_COEFF_FIELD_NUMBER: builtins.int
    MDL_MIN_WEIGHT_FIELD_NUMBER: builtins.int
    BENEFIT_REVISIT_SCALE_FIELD_NUMBER: builtins.int
    MAX_EVENT_BENEFIT_FIELD_NUMBER: builtins.int
    MAX_TOTAL_BENEFIT_FIELD_NUMBER: builtins.int
    MDL_HARD_LIMIT_FIELD_NUMBER: builtins.int
    HARD_LIMIT_MIN_BENEFIT_FIELD_NUMBER: builtins.int
    MDL_REGULARIZE_FIELD_NUMBER: builtins.int
    l2: builtins.float
    lr_power: builtins.float
    min_servable_mdl_benefit: builtins.float
    mdl_mix_in_margin: builtins.float
    mdl_benefit_rampup_coeff: builtins.float
    mdl_min_weight: builtins.float
    benefit_revisit_scale: builtins.float
    max_event_benefit: builtins.float
    max_total_benefit: builtins.float
    mdl_hard_limit: builtins.float
    hard_limit_min_benefit: builtins.bool
    mdl_regularize: builtins.bool
    def __init__(
        self,
        *,
        l2: builtins.float | None = ...,
        lr_power: builtins.float | None = ...,
        min_servable_mdl_benefit: builtins.float | None = ...,
        mdl_mix_in_margin: builtins.float | None = ...,
        mdl_benefit_rampup_coeff: builtins.float | None = ...,
        mdl_min_weight: builtins.float | None = ...,
        benefit_revisit_scale: builtins.float | None = ...,
        max_event_benefit: builtins.float | None = ...,
        max_total_benefit: builtins.float | None = ...,
        mdl_hard_limit: builtins.float | None = ...,
        hard_limit_min_benefit: builtins.bool | None = ...,
        mdl_regularize: builtins.bool | None = ...,
    ) -> None: ...
    def ClearField(
        self,
        field_name: typing.Literal[
            "benefit_revisit_scale",
            b"benefit_revisit_scale",
            "hard_limit_min_benefit",
            b"hard_limit_min_benefit",
            "l2",
            b"l2",
            "lr_power",
            b"lr_power",
            "max_event_benefit",
            b"max_event_benefit",
            "max_total_benefit",
            b"max_total_benefit",
            "mdl_benefit_rampup_coeff",
            b"mdl_benefit_rampup_coeff",
            "mdl_hard_limit",
            b"mdl_hard_limit",
            "mdl_min_weight",
            b"mdl_min_weight",
            "mdl_mix_in_margin",
            b"mdl_mix_in_margin",
            "mdl_regularize",
            b"mdl_regularize",
            "min_servable_mdl_benefit",
            b"min_servable_mdl_benefit",
        ],
    ) -> None: ...

global___MdlAdagradLightParameters = MdlAdagradLightParameters

@typing.final
class AdadeltaParameters(google.protobuf.message.Message):
    """https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adadelta
    https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L933
    """

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    RHO_FIELD_NUMBER: builtins.int
    EPSILON_FIELD_NUMBER: builtins.int
    rho: builtins.float
    epsilon: builtins.float
    def __init__(self, *, rho: builtins.float | None = ..., epsilon: builtins.float | None = ...) -> None: ...
    def ClearField(self, field_name: typing.Literal["epsilon", b"epsilon", "rho", b"rho"]) -> None: ...

global___AdadeltaParameters = AdadeltaParameters

@typing.final
class ProximalAdagradParameters(google.protobuf.message.Message):
    """https://www.tensorflow.org/api_docs/python/tf/compat/v1/train/ProximalAdagradOptimizer
    https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L1961
    """

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    L1_FIELD_NUMBER: builtins.int
    L2_FIELD_NUMBER: builtins.int
    l1: builtins.float
    l2: builtins.float
    def __init__(self, *, l1: builtins.float | None = ..., l2: builtins.float | None = ...) -> None: ...
    def ClearField(self, field_name: typing.Literal["l1", b"l1", "l2", b"l2"]) -> None: ...

global___ProximalAdagradParameters = ProximalAdagradParameters

@typing.final
class OnlineYogiParameters(google.protobuf.message.Message):
    """The online Yogi optimizer does not implement hyper-parameter update; use the
    dynamic learning rate feature instead, setting the learning rate to:
    user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
    Here, t is the current timestep.

    https://papers.nips.cc/paper/8186-adaptive-methods-for-nonconvex-optimization.pdf
    plus some extensions based on FTRL.

    Note that the code by default implements the lazy version of online Yogi.
    """

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    L1_FIELD_NUMBER: builtins.int
    L2_FIELD_NUMBER: builtins.int
    BETA2_FIELD_NUMBER: builtins.int
    l1: builtins.float
    """The L1 regularization parameter (used analogously to the one in FTRL)."""
    l2: builtins.float
    """The L2 regularization parameter (used analogously to the one in FTRL)."""
    beta2: builtins.float
    """\\beta_2 from Algorithm 2 in the paper."""
    def __init__(
        self, *, l1: builtins.float | None = ..., l2: builtins.float | None = ..., beta2: builtins.float | None = ...
    ) -> None: ...
    def ClearField(self, field_name: typing.Literal["beta2", b"beta2", "l1", b"l1", "l2", b"l2"]) -> None: ...

global___OnlineYogiParameters = OnlineYogiParameters

@typing.final
class ProximalYogiParameters(google.protobuf.message.Message):
    """The online Yogi optimizer does not implement hyper-parameter update; use the
    dynamic learning rate feature instead, setting the learning rate to:
    user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
    Here, t is the current timestep.

    https://papers.nips.cc/paper/8186-adaptive-methods-for-nonconvex-optimization.pdf
    plus some extensions based on FTRL.

    Note that the code by default implements the lazy version of proximal Yogi.
    """

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    L1_FIELD_NUMBER: builtins.int
    L2_FIELD_NUMBER: builtins.int
    BETA1_FIELD_NUMBER: builtins.int
    BETA2_FIELD_NUMBER: builtins.int
    EPSILON_FIELD_NUMBER: builtins.int
    l1: builtins.float
    """The L1 regularization parameter."""
    l2: builtins.float
    """The L2 regularization parameter."""
    beta1: builtins.float
    """The exponential decay rate for the 1st moment estimates."""
    beta2: builtins.float
    """The exponential decay rate for the 2nd moment estimates."""
    epsilon: builtins.float
    """A constant trading off adaptivity and noise."""
    def __init__(
        self,
        *,
        l1: builtins.float | None = ...,
        l2: builtins.float | None = ...,
        beta1: builtins.float | None = ...,
        beta2: builtins.float | None = ...,
        epsilon: builtins.float | None = ...,
    ) -> None: ...
    def ClearField(
        self, field_name: typing.Literal["beta1", b"beta1", "beta2", b"beta2", "epsilon", b"epsilon", "l1", b"l1", "l2", b"l2"]
    ) -> None: ...

global___ProximalYogiParameters = ProximalYogiParameters

@typing.final
class FrequencyEstimatorParameters(google.protobuf.message.Message):
    """Estimator for the frequency of updates to a lookup table. It maintains an
    array (tf.Variable) D, where each element records the average number of
    global steps between two consecutive batches that hit the corresponding
    bucket. Once an item with bucket id i is sampled, D[i] is updated by:
      D[i] <- D[i] * (1 - tau) + delta[i] * tau,

    where tau is a learning rate between 0 and 1 (exclusive), and
      delta[i] = current global step - last step i is sampled.

    The estimated frequency (sampling rate in a batch) is thus 1 / D[i].

    Elements in D are initialized with a large value max_delta. delta[i] will
    also be capped by this value.

    The exact sequence of operations used in the optimizer is shown below.
    last_hit_step[i] is a tf.Variable that holds the last global step at which i
    was sampled.

      delta = global_step - last_hit_step[i]
      clipped_delta = min(delta, params.max_delta)
      is_outlier = (delta >= params.outlier_threshold * D[i])
      D[i] <- is_outlier ? clipped_delta
                         : D[i] * (1 - params.tau) + clipped_delta * params.tau
      last_hit_step[i] <- global_step
    """

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    TAU_FIELD_NUMBER: builtins.int
    MAX_DELTA_FIELD_NUMBER: builtins.int
    OUTLIER_THRESHOLD_FIELD_NUMBER: builtins.int
    WEIGHT_EXPONENT_FIELD_NUMBER: builtins.int
    tau: builtins.float
    """Learning rate between (0, 1) that is used to update the array D."""
    max_delta: builtins.float
    """Maximum value of delta: difference between the current global step and the
    last global step at which the row was sampled.
    """
    outlier_threshold: builtins.float
    """Threshold used to determine whether the current update is an outlier."""
    weight_exponent: builtins.float
    """The weight exponent used to transform the estimated delta into weights.
    The transformation function is: (delta / max_delta) ^ (weight_exponent)
    """
    def __init__(
        self,
        *,
        tau: builtins.float | None = ...,
        max_delta: builtins.float | None = ...,
        outlier_threshold: builtins.float | None = ...,
        weight_exponent: builtins.float | None = ...,
    ) -> None: ...
    def ClearField(
        self,
        field_name: typing.Literal[
            "max_delta",
            b"max_delta",
            "outlier_threshold",
            b"outlier_threshold",
            "tau",
            b"tau",
            "weight_exponent",
            b"weight_exponent",
        ],
    ) -> None: ...

global___FrequencyEstimatorParameters = FrequencyEstimatorParameters

@typing.final
class UserDefinedProgramParameters(google.protobuf.message.Message):
    """A user-defined optimizer.
    The contained HLO program must take the following arguments in the following
    order:
    1.  gradients
    2.  table weights
    3.  slot variables
    4.  an optional scalar input that is passed in via the dynamic learning
        rate mechanism.

    It must return/end in a tuple op that contains the following values in the
    following order:
    1.  new table values
    2.  new slot variable value

    The program must have shape (1,1) with dtype float32 throughout and only use
    HLO that operate elementwise (e.g., no reduce, no variables, no control flow
    and no broadcasting outside of the single scalar input).
    The HLO program should be written as if it were a dense update. It will be
    called on each row that needs an update and will applied elementwise.
    """

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    PROGRAM_FIELD_NUMBER: builtins.int
    @property
    def program(self) -> tensorflow.compiler.xla.service.hlo_pb2.HloModuleProto: ...
    def __init__(self, *, program: tensorflow.compiler.xla.service.hlo_pb2.HloModuleProto | None = ...) -> None: ...
    def HasField(self, field_name: typing.Literal["program", b"program"]) -> builtins.bool: ...
    def ClearField(self, field_name: typing.Literal["program", b"program"]) -> None: ...

global___UserDefinedProgramParameters = UserDefinedProgramParameters

@typing.final
class AssignParameters(google.protobuf.message.Message):
    """Optimizer that just sets the variable to the value of the gradient. To be
    correct, this requires either gradient accumulation (to sum the values of a
    computed expression across the samples) or to deduplicate IDs within a single
    host (to assign the value from an arbitrary sample).
    """

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    def __init__(self) -> None: ...

global___AssignParameters = AssignParameters

@typing.final
class GradientAccumulationStatus(google.protobuf.message.Message):
    """Status of using gradient accumulation (doing two passes over the input
    gradients: one to accumulate them into a temporary array and another to apply
    them using the actual optimization algorithm). The extra message is to wrap
    the enum for scoping.
    """

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    class _Status:
        ValueType = typing.NewType("ValueType", builtins.int)
        V: typing_extensions.TypeAlias = ValueType

    class _StatusEnumTypeWrapper(
        google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[GradientAccumulationStatus._Status.ValueType], builtins.type
    ):
        DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
        UNSPECIFIED: GradientAccumulationStatus._Status.ValueType  # 0
        ENABLED: GradientAccumulationStatus._Status.ValueType  # 1
        DISABLED: GradientAccumulationStatus._Status.ValueType  # 2

    class Status(_Status, metaclass=_StatusEnumTypeWrapper):
        """if UNSPECIFIED (default), gradient accumulation is ENABLED."""

    UNSPECIFIED: GradientAccumulationStatus.Status.ValueType  # 0
    ENABLED: GradientAccumulationStatus.Status.ValueType  # 1
    DISABLED: GradientAccumulationStatus.Status.ValueType  # 2

    def __init__(self) -> None: ...

global___GradientAccumulationStatus = GradientAccumulationStatus

@typing.final
class LowDimensionalPackingStatus(google.protobuf.message.Message):
    """Whether to optimize the packing of low-dimensional embedding tables in HBM
    (high bandwidth memory). TPUs access HBM at 32-byte (8-float) granularity.
    For functional correctness, the TPU software internally pads the embedding
    dimension to a multiple of 8. This can sometimes lead to significant memory
    wastage due to padding. For 1-dimensional, 2-dimensional, and 4-dimensional,
    the TPU software can remove this padding by packing multiple rows into the
    same 8-float HBM chunk. For example, 8 rows could be packed into the same
    8-float chunk for a 1-dimensional embedding table.

    There is one important limitation for this HBM packing though. When only a
    subset of rows in an 8-float chunk are accessed on a particular step, the
    adjoining rows in the same chunk are updated with zero gradients on the
    backward pass even if they are not touched. This is an artifact of the
    packing implementation. This operation is NOT functionally correct for
    optimizers where zero gradients change the embeddings/slot-variable values,
    e.g., momentum-based optimizers. Hence, this HBM packing cannot be enabled
    for embedding tables with such optimizers. The TPU software automatically
    recognizes that a zero gradient can modify state and turns off the low
    dimensional embedding packing in that scenario.

    However, for optimizers where a zero gradient is a NoOp, such as SGD,
    Adagrad, and FTRL, this packing optimization can be used. However, there are
    some important considerations:
    * Clipping limits: The initial values for such embeddings should fall within
      the clipping limits specified in the optimization parameters. Otherwise, a
      zero gradient will cause the embeddings to be clipped. This changes state
      and hence, is not a NoOp.
    * FTRL: The embedding vector is computed directly from the values of the
      accumulator and linear slot variables. Hence, the initial embedding values
      should match that computed from the initial values of the accumulator and
      linear slot variables. Note that in nearly all cases, the linear value is
      initialized to zero; this corresponds to an embedding value of zero.

    Performance: The TPU has to perform additional work when low dimensional
    packing is enabled. In certain situations when the vocabulary size is small,
    it may not make sense to turn on this packing since the total memory usage
    due to padding is extremely low. Hence, the TPU software automatically turns
    off the packing optimization in such scenarios.
    """

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    class _Status:
        ValueType = typing.NewType("ValueType", builtins.int)
        V: typing_extensions.TypeAlias = ValueType

    class _StatusEnumTypeWrapper(
        google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[LowDimensionalPackingStatus._Status.ValueType], builtins.type
    ):
        DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
        UNSPECIFIED: LowDimensionalPackingStatus._Status.ValueType  # 0
        ENABLED: LowDimensionalPackingStatus._Status.ValueType  # 1
        DISABLED: LowDimensionalPackingStatus._Status.ValueType  # 2

    class Status(_Status, metaclass=_StatusEnumTypeWrapper):
        """if UNSPECIFIED (default), the low dimension packing status is DISABLED.
        This can change in future.

        if ENABLED, the low dimension packing is enabled only if the following
        three additional conditions are true:
         * The optimizer treats the zero gradient as a NoOp.
         * The embedding dimension is 1, 2, or 4.
         * The vocabulary size is large enough to avoid performance issues.

        if DISABLED, the low dimension packing is always disabled.
        """

    UNSPECIFIED: LowDimensionalPackingStatus.Status.ValueType  # 0
    ENABLED: LowDimensionalPackingStatus.Status.ValueType  # 1
    DISABLED: LowDimensionalPackingStatus.Status.ValueType  # 2

    def __init__(self) -> None: ...

global___LowDimensionalPackingStatus = LowDimensionalPackingStatus

@typing.final
class HotIdReplicationConfiguration(google.protobuf.message.Message):
    """Configuration proto for hot ID optimization. This is an experimental feature
    that is currently disabled (by default).
    """

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    class _Status:
        ValueType = typing.NewType("ValueType", builtins.int)
        V: typing_extensions.TypeAlias = ValueType

    class _StatusEnumTypeWrapper(
        google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[HotIdReplicationConfiguration._Status.ValueType],
        builtins.type,
    ):
        DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
        UNSPECIFIED: HotIdReplicationConfiguration._Status.ValueType  # 0
        ENABLED: HotIdReplicationConfiguration._Status.ValueType  # 1
        DISABLED: HotIdReplicationConfiguration._Status.ValueType  # 2
        MIGRATION_ONLY: HotIdReplicationConfiguration._Status.ValueType  # 3

    class Status(_Status, metaclass=_StatusEnumTypeWrapper):
        """Whether to enable or disable hot ID optimization.
        If set to UNSPECIFIED (default), hot ID optimization is DISABLED.
        If set to ENABLED, hot ID replication is turned ON.
        If set to MIGRATION_ONLY, hot ID migration is turned ON.
        """

    UNSPECIFIED: HotIdReplicationConfiguration.Status.ValueType  # 0
    ENABLED: HotIdReplicationConfiguration.Status.ValueType  # 1
    DISABLED: HotIdReplicationConfiguration.Status.ValueType  # 2
    MIGRATION_ONLY: HotIdReplicationConfiguration.Status.ValueType  # 3

    STATUS_FIELD_NUMBER: builtins.int
    status: global___HotIdReplicationConfiguration.Status.ValueType
    def __init__(self, *, status: global___HotIdReplicationConfiguration.Status.ValueType | None = ...) -> None: ...
    def ClearField(self, field_name: typing.Literal["status", b"status"]) -> None: ...

global___HotIdReplicationConfiguration = HotIdReplicationConfiguration

@typing.final
class OptimizationParameters(google.protobuf.message.Message):
    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    LEARNING_RATE_FIELD_NUMBER: builtins.int
    CLIPPING_LIMITS_FIELD_NUMBER: builtins.int
    GRADIENT_CLIPPING_LIMITS_FIELD_NUMBER: builtins.int
    WEIGHT_DECAY_FACTOR_FIELD_NUMBER: builtins.int
    MULTIPLY_WEIGHT_DECAY_FACTOR_BY_LEARNING_RATE_FIELD_NUMBER: builtins.int
    SIMULATED_QUANTIZATION_FIELD_NUMBER: builtins.int
    GRADIENT_ACCUMULATION_STATUS_FIELD_NUMBER: builtins.int
    LOW_DIMENSIONAL_PACKING_STATUS_FIELD_NUMBER: builtins.int
    HOT_ID_REPLICATION_CONFIGURATION_FIELD_NUMBER: builtins.int
    ADAGRAD_FIELD_NUMBER: builtins.int
    ADAGRAD_MOMENTUM_FIELD_NUMBER: builtins.int
    BOUNDED_ADAGRAD_FIELD_NUMBER: builtins.int
    STOCHASTIC_GRADIENT_DESCENT_FIELD_NUMBER: builtins.int
    FTRL_FIELD_NUMBER: builtins.int
    ADAM_FIELD_NUMBER: builtins.int
    MOMENTUM_FIELD_NUMBER: builtins.int
    LION_FIELD_NUMBER: builtins.int
    RMS_PROP_FIELD_NUMBER: builtins.int
    CENTERED_RMS_PROP_FIELD_NUMBER: builtins.int
    MDL_ADAGRAD_LIGHT_FIELD_NUMBER: builtins.int
    ADADELTA_FIELD_NUMBER: builtins.int
    PROXIMAL_ADAGRAD_FIELD_NUMBER: builtins.int
    ONLINE_YOGI_FIELD_NUMBER: builtins.int
    PROXIMAL_YOGI_FIELD_NUMBER: builtins.int
    FREQUENCY_ESTIMATOR_FIELD_NUMBER: builtins.int
    USER_DEFINED_PROGRAM_FIELD_NUMBER: builtins.int
    ASSIGN_FIELD_NUMBER: builtins.int
    weight_decay_factor: builtins.float
    """Amount of weight decay to apply; see weight_decay_optimizers.py for
    details. All optimizers except MDL Adagrad Light are supported with this
    option. Although there is no check, users who want weight decay will also
    want to ensure that gradient accumulation is enabled so that the decay will
    happen once per global batch.
    """
    multiply_weight_decay_factor_by_learning_rate: builtins.bool
    """If true, the weight decay factor is multiplied by the current learning rate
    before use; this is to match the note in DecoupledWeightDecayExtension in
    weight_decay_optimizers.py.
    """
    gradient_accumulation_status: global___GradientAccumulationStatus.Status.ValueType
    """Status of using gradient accumulation (doing two passes over the input
    gradients: one to accumulate them into a temporary array and another to
    apply them using the actual optimization algorithm).
    """
    low_dimensional_packing_status: global___LowDimensionalPackingStatus.Status.ValueType
    """Status of the low-dimensional embedding packing optimization. This controls
    whether to optimize the packing of 1-dimensional, 2-dimensional, and
    4-dimensional embedding tables in memory.
    """
    @property
    def learning_rate(self) -> global___LearningRate:
        """Learning rate used for updating the embedding layer parameters."""

    @property
    def clipping_limits(self) -> global___ClippingLimits:
        """Limits to which to clip the weight values after the backward pass; not
        present means no limits are applied.
        """

    @property
    def gradient_clipping_limits(self) -> global___ClippingLimits:
        """Limits to which to clip the backward pass gradient before using it for
        updates; not present means no limits are applied.
        """

    @property
    def simulated_quantization(self) -> global___SimulatedQuantization:
        """Configuration for simulated quantization which is used to reduce
        training/serving skew when the serving variables are quantized. The same
        quantization operations are executed during training to minimize
        differences with serving.
        """

    @property
    def hot_id_replication_configuration(self) -> global___HotIdReplicationConfiguration:
        """Configuration proto for hot ID replication. This is an experimental
        feature that is currently disabled (by default).
        """

    @property
    def adagrad(self) -> global___AdagradParameters: ...
    @property
    def adagrad_momentum(self) -> global___AdagradMomentumParameters: ...
    @property
    def bounded_adagrad(self) -> global___BoundedAdagradParameters: ...
    @property
    def stochastic_gradient_descent(self) -> global___StochasticGradientDescentParameters: ...
    @property
    def ftrl(self) -> global___FtrlParameters: ...
    @property
    def adam(self) -> global___AdamParameters: ...
    @property
    def momentum(self) -> global___MomentumParameters: ...
    @property
    def lion(self) -> global___LionParameters: ...
    @property
    def rms_prop(self) -> global___RmsPropParameters: ...
    @property
    def centered_rms_prop(self) -> global___CenteredRmsPropParameters: ...
    @property
    def mdl_adagrad_light(self) -> global___MdlAdagradLightParameters: ...
    @property
    def adadelta(self) -> global___AdadeltaParameters: ...
    @property
    def proximal_adagrad(self) -> global___ProximalAdagradParameters: ...
    @property
    def online_yogi(self) -> global___OnlineYogiParameters: ...
    @property
    def proximal_yogi(self) -> global___ProximalYogiParameters: ...
    @property
    def frequency_estimator(self) -> global___FrequencyEstimatorParameters: ...
    @property
    def user_defined_program(self) -> global___UserDefinedProgramParameters: ...
    @property
    def assign(self) -> global___AssignParameters: ...
    def __init__(
        self,
        *,
        learning_rate: global___LearningRate | None = ...,
        clipping_limits: global___ClippingLimits | None = ...,
        gradient_clipping_limits: global___ClippingLimits | None = ...,
        weight_decay_factor: builtins.float | None = ...,
        multiply_weight_decay_factor_by_learning_rate: builtins.bool | None = ...,
        simulated_quantization: global___SimulatedQuantization | None = ...,
        gradient_accumulation_status: global___GradientAccumulationStatus.Status.ValueType | None = ...,
        low_dimensional_packing_status: global___LowDimensionalPackingStatus.Status.ValueType | None = ...,
        hot_id_replication_configuration: global___HotIdReplicationConfiguration | None = ...,
        adagrad: global___AdagradParameters | None = ...,
        adagrad_momentum: global___AdagradMomentumParameters | None = ...,
        bounded_adagrad: global___BoundedAdagradParameters | None = ...,
        stochastic_gradient_descent: global___StochasticGradientDescentParameters | None = ...,
        ftrl: global___FtrlParameters | None = ...,
        adam: global___AdamParameters | None = ...,
        momentum: global___MomentumParameters | None = ...,
        lion: global___LionParameters | None = ...,
        rms_prop: global___RmsPropParameters | None = ...,
        centered_rms_prop: global___CenteredRmsPropParameters | None = ...,
        mdl_adagrad_light: global___MdlAdagradLightParameters | None = ...,
        adadelta: global___AdadeltaParameters | None = ...,
        proximal_adagrad: global___ProximalAdagradParameters | None = ...,
        online_yogi: global___OnlineYogiParameters | None = ...,
        proximal_yogi: global___ProximalYogiParameters | None = ...,
        frequency_estimator: global___FrequencyEstimatorParameters | None = ...,
        user_defined_program: global___UserDefinedProgramParameters | None = ...,
        assign: global___AssignParameters | None = ...,
    ) -> None: ...
    def HasField(
        self,
        field_name: typing.Literal[
            "adadelta",
            b"adadelta",
            "adagrad",
            b"adagrad",
            "adagrad_momentum",
            b"adagrad_momentum",
            "adam",
            b"adam",
            "assign",
            b"assign",
            "bounded_adagrad",
            b"bounded_adagrad",
            "centered_rms_prop",
            b"centered_rms_prop",
            "clipping_limits",
            b"clipping_limits",
            "frequency_estimator",
            b"frequency_estimator",
            "ftrl",
            b"ftrl",
            "gradient_clipping_limits",
            b"gradient_clipping_limits",
            "hot_id_replication_configuration",
            b"hot_id_replication_configuration",
            "learning_rate",
            b"learning_rate",
            "lion",
            b"lion",
            "mdl_adagrad_light",
            b"mdl_adagrad_light",
            "momentum",
            b"momentum",
            "online_yogi",
            b"online_yogi",
            "parameters",
            b"parameters",
            "proximal_adagrad",
            b"proximal_adagrad",
            "proximal_yogi",
            b"proximal_yogi",
            "rms_prop",
            b"rms_prop",
            "simulated_quantization",
            b"simulated_quantization",
            "stochastic_gradient_descent",
            b"stochastic_gradient_descent",
            "user_defined_program",
            b"user_defined_program",
        ],
    ) -> builtins.bool: ...
    def ClearField(
        self,
        field_name: typing.Literal[
            "adadelta",
            b"adadelta",
            "adagrad",
            b"adagrad",
            "adagrad_momentum",
            b"adagrad_momentum",
            "adam",
            b"adam",
            "assign",
            b"assign",
            "bounded_adagrad",
            b"bounded_adagrad",
            "centered_rms_prop",
            b"centered_rms_prop",
            "clipping_limits",
            b"clipping_limits",
            "frequency_estimator",
            b"frequency_estimator",
            "ftrl",
            b"ftrl",
            "gradient_accumulation_status",
            b"gradient_accumulation_status",
            "gradient_clipping_limits",
            b"gradient_clipping_limits",
            "hot_id_replication_configuration",
            b"hot_id_replication_configuration",
            "learning_rate",
            b"learning_rate",
            "lion",
            b"lion",
            "low_dimensional_packing_status",
            b"low_dimensional_packing_status",
            "mdl_adagrad_light",
            b"mdl_adagrad_light",
            "momentum",
            b"momentum",
            "multiply_weight_decay_factor_by_learning_rate",
            b"multiply_weight_decay_factor_by_learning_rate",
            "online_yogi",
            b"online_yogi",
            "parameters",
            b"parameters",
            "proximal_adagrad",
            b"proximal_adagrad",
            "proximal_yogi",
            b"proximal_yogi",
            "rms_prop",
            b"rms_prop",
            "simulated_quantization",
            b"simulated_quantization",
            "stochastic_gradient_descent",
            b"stochastic_gradient_descent",
            "user_defined_program",
            b"user_defined_program",
            "weight_decay_factor",
            b"weight_decay_factor",
        ],
    ) -> None: ...
    def WhichOneof(
        self, oneof_group: typing.Literal["parameters", b"parameters"]
    ) -> (
        typing.Literal[
            "adagrad",
            "adagrad_momentum",
            "bounded_adagrad",
            "stochastic_gradient_descent",
            "ftrl",
            "adam",
            "momentum",
            "lion",
            "rms_prop",
            "centered_rms_prop",
            "mdl_adagrad_light",
            "adadelta",
            "proximal_adagrad",
            "online_yogi",
            "proximal_yogi",
            "frequency_estimator",
            "user_defined_program",
            "assign",
        ]
        | None
    ): ...

global___OptimizationParameters = OptimizationParameters

@typing.final
class StateVariableSpecification(google.protobuf.message.Message):
    """Specification of an optimization algorithm's state variables (both the main
    value vector and any extra accumulators, etc.). This proto is only used
    internally by the TPU software and is not exposed directly to the TF model.
    """

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    @typing.final
    class UserDefined(google.protobuf.message.Message):
        """A normal state variable that should be saved and restored in checkpoints
        and used as an input or output to non-debug TensorFlow ops.
        """

        DESCRIPTOR: google.protobuf.descriptor.Descriptor

        def __init__(self) -> None: ...

    @typing.final
    class FillWithConstant(google.protobuf.message.Message):
        """A state variable that should be filled with a constant and normally hidden
        from users (used for intermediate gradients being accumulated, for
        example).
        """

        DESCRIPTOR: google.protobuf.descriptor.Descriptor

        INITIAL_VALUE_FIELD_NUMBER: builtins.int
        initial_value: builtins.float
        def __init__(self, *, initial_value: builtins.float | None = ...) -> None: ...
        def ClearField(self, field_name: typing.Literal["initial_value", b"initial_value"]) -> None: ...

    NAME_FIELD_NUMBER: builtins.int
    USER_DEFINED_FIELD_NUMBER: builtins.int
    FILL_WITH_CONSTANT_FIELD_NUMBER: builtins.int
    name: builtins.str
    """Parameter name for the state variable."""
    @property
    def user_defined(self) -> global___StateVariableSpecification.UserDefined: ...
    @property
    def fill_with_constant(self) -> global___StateVariableSpecification.FillWithConstant: ...
    def __init__(
        self,
        *,
        name: builtins.str | None = ...,
        user_defined: global___StateVariableSpecification.UserDefined | None = ...,
        fill_with_constant: global___StateVariableSpecification.FillWithConstant | None = ...,
    ) -> None: ...
    def HasField(
        self,
        field_name: typing.Literal[
            "fill_with_constant", b"fill_with_constant", "usage", b"usage", "user_defined", b"user_defined"
        ],
    ) -> builtins.bool: ...
    def ClearField(
        self,
        field_name: typing.Literal[
            "fill_with_constant", b"fill_with_constant", "name", b"name", "usage", b"usage", "user_defined", b"user_defined"
        ],
    ) -> None: ...
    def WhichOneof(
        self, oneof_group: typing.Literal["usage", b"usage"]
    ) -> typing.Literal["user_defined", "fill_with_constant"] | None: ...

global___StateVariableSpecification = StateVariableSpecification
