text: # Copyright 2021 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """RMSprop optimizer implementation."""
  
  from keras.optimizer_experimental import optimizer
  from keras.utils import generic_utils
  import tensorflow.compat.v2 as tf
  
  
  @generic_utils.register_keras_serializable()
  class RMSprop(optimizer.Optimizer):
    r"""Optimizer that implements the RMSprop algorithm.
  
    The gist of RMSprop is to:
  
    - Maintain a moving (discounted) average of the square of gradients
    - Divide the gradient by the root of this average
  
    This implementation of RMSprop uses plain momentum, not Nesterov momentum.
  
    The centered version additionally maintains a moving average of the
    gradients, and uses that average to estimate the variance.
  
    Attributes:
      learning_rate: Initial value for the learning rate:
        either a floating point value,
        or a `tf.keras.optimizers.schedules.LearningRateSchedule` instance.
        Defaults to 0.001.
      rho: float, defaults to 0.9. Discounting factor for the old gradients.
      momentum: float, defaults to 0.0. If not 0.0., the optimizer tracks the
        momentum value, with a decay rate equals to `1 - momentum`.
      epsilon: A small constant for numerical stability. This epsilon is
        "epsilon hat" in the Kingma and Ba paper (in the formula just before
        Section 2.1), not the epsilon in Algorithm 1 of the paper. Defaults to
        1e-7.
      centered: Boolean. If `True`, gradients are normalized by the estimated
        variance of the gradient; if False, by the uncentered second moment.
        Setting this to `True` may help with training, but is slightly more
        expensive in terms of computation and memory. Defaults to `False`.
      gradients_clip_option: an instance of
        `optimizer_experimental.GradientsClipOption`, for attributes related to
        gradients clipping, such as clipnorm and clipvalue. Default to None
        (not applying gradients clipping).
      ema_option: an instance of `optimizer_experimental.EMAOption`, for
        attributes related to exponenatial moving average, such as `use_ema` (a
        boolean field indicates if EMA is used) and EMA momentum. Default to None
        (not applying EMA).
      jit_compile: Bool, default to False. If True, the optimizer will use XLA
        acceleration. `jit_compile` can only be False when using Parameter
        Server Strategy.
      name: Optional name prefix for the operations created when applying
        gradients.  Defaults to `"RMSprop"`.
      **kwargs: keyword arguments only used for backward compatibility with
        `optimizer_v2.OptimizerV2`. Any new code using
        `optimizer_experimental.Optimizer` should leave this parameter empty.
  
    Usage:
  
    >>> opt = tf.keras.optimizers.RMSprop(learning_rate=0.1)
    >>> var1 = tf.Variable(10.0)
    >>> loss = lambda: (var1 ** 2) / 2.0    # d(loss) / d(var1) = var1
    >>> step_count = opt.minimize(loss, [var1]).numpy()
    >>> var1.numpy()
    9.683772
  
    Reference:
      - [Hinton, 2012](
        http://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf)
  
    """
  
    def __init__(self,
                 learning_rate=0.001,
                 rho=0.9,
                 momentum=0.0,
                 epsilon=1e-7,
                 centered=False,
                 gradients_clip_option=None,
                 ema_option=None,
                 jit_compile=False,
                 name='RMSprop',
                 **kwargs):
      super(RMSprop, self).__init__(
          gradients_clip_option=gradients_clip_option,
          ema_option=ema_option,
          jit_compile=jit_compile,
          name=name,
          **kwargs)
      self._learning_rate = self._build_learning_rate(learning_rate)
      self.rho = rho
      self.momentum = momentum
      self.epsilon = epsilon
      self.centered = centered
  
    def build(self, var_list):
      super().build(var_list)
      if hasattr(self, '_built') and self._built:
        return
      self._built = True
  
      self._velocities = []
      for var in var_list:
        self._velocities.append(
            self.add_variable_from_reference(var, 'velocity'))
  
      self._momentums = []
      if self.momentum > 0:
        for var in var_list:
          self._momentums.append(
              self.add_variable_from_reference(var, 'momentum'))
  
      self._average_gradients = []
      if self.centered:
        for var in var_list:
          self._average_gradients.append(
              self.add_variable_from_reference(var, 'average_gradient'))
  
    def update_step(self, gradient, variable, params=None):
      """Update step given gradient and the associated model variable."""
      if self._var_key(variable) not in self._index_dict:
        raise KeyError(f'Optimizer cannot recognize variable {variable.name}, '
                       f'this usually means you are calling an optimizer '
                       f'previously used on a different model. Please try '
                       f'creating a new optimizer instance.')
      lr = tf.cast(self.learning_rate, variable.dtype)
  
      var_key = self._var_key(variable)
      velocity = self._velocities[self._index_dict[var_key]]
      momentum = None
      if self.momentum > 0:
        momentum = self._momentums[self._index_dict[var_key]]
      average_grad = None
      if self.centered:
        average_grad = self._average_gradients[self._index_dict[var_key]]
  
      rho = self.rho
  
      if isinstance(gradient, tf.IndexedSlices):
        # Sparse gradients.
        velocity.assign(rho * velocity)
        velocity.scatter_add(tf.IndexedSlices(
            tf.square(gradient.values) * (1 - rho), gradient.indices))
        if self.centered:
          average_grad.assign(rho * average_grad)
          average_grad.scatter_add(
              tf.IndexedSlices(
                  tf.square(gradient.values) * (1 - rho), gradient.indices))
          velocity.assign_add(-tf.square(average_grad))
        velocity_value = tf.gather(velocity, gradient.indices)
        transformed_grad = tf.IndexedSlices(
            gradient.values / (tf.sqrt(velocity_value) + self.epsilon),
            gradient.indices)
  
        if self.momentum > 0:
          momentum.assign(self.momentum * momentum)
          momentum.scatter_add(transformed_grad)
          variable.assign_add(-lr * momentum)
        else:
          variable.scatter_add(
              tf.IndexedSlices(-lr * transformed_grad.values,
                               transformed_grad.indices))
      else:
        # Dense gradients.
        velocity.assign(rho * velocity + (1 - rho) * tf.square(gradient))
        if self.centered:
          average_grad.assign(rho * average_grad +
                              (1 - rho) * tf.square(gradient))
          velocity.assign_add(-tf.square(average_grad))
        transformed_grad = gradient / (tf.sqrt(velocity) + self.epsilon)
        if self.momentum > 0:
          momentum.assign(self.momentum * momentum + transformed_grad)
          variable.assign_add(-lr * momentum)
        else:
          variable.assign_add(-lr * transformed_grad)
  
    def get_config(self):
      config = super(RMSprop, self).get_config()
  
      config.update({
          'learning_rate': self._serialize_hyperparameter(self._learning_rate),
          'rho': self.rho,
          'momentum': self.momentum,
          'epsilon': self.epsilon,
          'centered': self.centered,
      })
      return config
  