basedpyright 1.36.2 → 1.37.0-961d59bacff1958014fe4e9ed2ef3116fe2266e2

package/dist/typeshed-fallback/stubs/tensorflow/tensorflow/core/protobuf/tpu/optimization_parameters_pb2.pyi

@@ -0,0 +1,1326 @@
+"""
+@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