chore: import upstream snapshot with attribution

tensorflow/python/keras/optimizer_v2/BUILD

@@ -0,0 +1,104 @@
+# Description:
+#   Contains the Keras OptimizerV2 API (internal TensorFlow version).
+
+load("@xla//third_party/rules_python/python:defs.bzl", "py_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    # TODO(scottzhu): Remove non-keras deps from TF.
+    default_visibility = [
+        "//tensorflow/python:__pkg__",
+        "//tensorflow/python/distribute:__pkg__",
+        "//tensorflow/python/keras:__subpackages__",
+        "//tensorflow/python/tpu/tests:__pkg__",
+        "//tensorflow/python/trackable:__pkg__",
+        "//tensorflow/python/training:__pkg__",
+    ],
+    licenses = ["notice"],
+)
+
+filegroup(
+    name = "all_py_srcs",
+    srcs = glob(["*.py"]),
+    visibility = ["//tensorflow/python/keras/google/private_tf_api_test:__pkg__"],
+)
+
+py_library(
+    name = "optimizer_v2",
+    srcs = [
+        "adadelta.py",
+        "adagrad.py",
+        "adam.py",
+        "adamax.py",
+        "ftrl.py",
+        "gradient_descent.py",
+        "nadam.py",
+        "optimizer_v2.py",
+        "rmsprop.py",
+        "utils.py",
+    ],
+    srcs_version = "PY3",
+    strict_deps = False,
+    deps = [
+        ":learning_rate_schedule",
+        "//tensorflow/python/distribute:central_storage_strategy",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/distribute:parameter_server_strategy",
+        "//tensorflow/python/distribute:parameter_server_strategy_v2",
+        "//tensorflow/python/distribute:reduce_util",
+        "//tensorflow/python/distribute:values",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:indexed_slices",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:tensor_conversion",
+        "//tensorflow/python/keras:backend",
+        "//tensorflow/python/keras:backend_config",
+        "//tensorflow/python/keras/engine:base_layer_utils",
+        "//tensorflow/python/keras/initializers",
+        "//tensorflow/python/keras/utils:layer_utils",
+        "//tensorflow/python/keras/utils:tf_utils",
+        "//tensorflow/python/ops:control_flow_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/ops:state_ops",
+        "//tensorflow/python/ops:variable_scope",
+        "//tensorflow/python/ops:variables",
+    ],
+)
+
+py_library(
+    name = "learning_rate_schedule",
+    srcs = [
+        "learning_rate_schedule.py",
+    ],
+    srcs_version = "PY3",
+    strict_deps = False,
+    deps = [
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_conversion",
+        "//tensorflow/python/keras/utils:generic_utils",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:control_flow_case",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:random_ops",
+    ],
+)
+
+py_library(
+    name = "legacy_learning_rate_decay",
+    srcs = ["legacy_learning_rate_decay.py"],
+    srcs_version = "PY3",
+    strict_deps = False,
+    deps = [
+        ":learning_rate_schedule",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_conversion",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/util:tf_export",
+    ],
+)

tensorflow/python/keras/optimizer_v2/adadelta.py

@@ -0,0 +1,149 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Adadelta optimizer implementation."""
+# pylint: disable=g-classes-have-attributes
+
+import numpy as np
+from tensorflow.python.framework import tensor_conversion
+from tensorflow.python.keras import backend_config
+from tensorflow.python.keras.optimizer_v2 import optimizer_v2
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_training_ops
+
+
+class Adadelta(optimizer_v2.OptimizerV2):
+  r"""Optimizer that implements the Adadelta algorithm.
+
+  Adadelta optimization is a stochastic gradient descent method that is based on
+  adaptive learning rate per dimension to address two drawbacks:
+
+  - The continual decay of learning rates throughout training.
+  - The need for a manually selected global learning rate.
+
+  Adadelta is a more robust extension of Adagrad that adapts learning rates
+  based on a moving window of gradient updates, instead of accumulating all
+  past gradients. This way, Adadelta continues learning even when many updates
+  have been done. Compared to Adagrad, in the original version of Adadelta you
+  don't have to set an initial learning rate. In this version, the initial
+  learning rate can be set, as in most other Keras optimizers.
+
+  Args:
+    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.
+      Note that `Adadelta` tends to benefit from higher initial learning rate
+      values compared to other optimizers.
+      To match the exact form in the original paper, use 1.0.
+    rho: A `Tensor` or a floating point value. The decay rate.
+    epsilon: Small floating point value used to maintain numerical stability.
+    name: Optional name prefix for the operations created when applying
+      gradients.  Defaults to `"Adadelta"`.
+    **kwargs: Keyword arguments. Allowed to be one of
+      `"clipnorm"` or `"clipvalue"`.
+      `"clipnorm"` (float) clips gradients by norm and represents
+      the maximum norm of each parameter;
+      `"clipvalue"` (float) clips gradient by value and represents the
+      maximum absolute value of each parameter.
+
+  Reference:
+    - [Zeiler, 2012](http://arxiv.org/abs/1212.5701)
+  """
+
+  _HAS_AGGREGATE_GRAD = True
+
+  def __init__(self,
+               learning_rate=0.001,
+               rho=0.95,
+               epsilon=1e-7,
+               name='Adadelta',
+               **kwargs):
+    super(Adadelta, self).__init__(name, **kwargs)
+    self._set_hyper('learning_rate', kwargs.get('lr', learning_rate))
+    self._set_hyper('decay', self._initial_decay)
+    self._set_hyper('rho', rho)
+    self.epsilon = epsilon or backend_config.epsilon()
+
+  def _create_slots(self, var_list):
+    # Separate for-loops to respect the ordering of slot variables from v1.
+    for v in var_list:
+      self.add_slot(v, 'accum_grad')
+    for v in var_list:
+      self.add_slot(v, 'accum_var')
+
+  def _prepare_local(self, var_device, var_dtype, apply_state):
+    super(Adadelta, self)._prepare_local(var_device, var_dtype, apply_state)
+    apply_state[(var_device, var_dtype)].update(
+        dict(
+            epsilon=tensor_conversion.convert_to_tensor_v2_with_dispatch(
+                self.epsilon, var_dtype
+            ),
+            rho=array_ops.identity(self._get_hyper('rho', var_dtype)),
+        )
+    )
+
+  def set_weights(self, weights):
+    params = self.weights
+    # Override set_weights for backward compatibility of Keras V1 optimizer
+    # since it does not include iteration at head of the weight list. Set
+    # iteration to 0.
+    if len(params) == len(weights) + 1:
+      weights = [np.array(0)] + weights
+    super(Adadelta, self).set_weights(weights)
+
+  def _resource_apply_dense(self, grad, var, apply_state=None):
+    var_device, var_dtype = var.device, var.dtype.base_dtype
+    coefficients = ((apply_state or {}).get((var_device, var_dtype))
+                    or self._fallback_apply_state(var_device, var_dtype))
+
+    accum_grad = self.get_slot(var, 'accum_grad')
+    accum_var = self.get_slot(var, 'accum_var')
+    return gen_training_ops.ResourceApplyAdadelta(
+        var=var.handle,
+        accum=accum_grad.handle,
+        accum_update=accum_var.handle,
+        lr=coefficients['lr_t'],
+        rho=coefficients['rho'],
+        epsilon=coefficients['epsilon'],
+        grad=grad,
+        use_locking=self._use_locking)
+
+  def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
+    var_device, var_dtype = var.device, var.dtype.base_dtype
+    coefficients = ((apply_state or {}).get((var_device, var_dtype))
+                    or self._fallback_apply_state(var_device, var_dtype))
+
+    accum_grad = self.get_slot(var, 'accum_grad')
+    accum_var = self.get_slot(var, 'accum_var')
+    return gen_training_ops.ResourceSparseApplyAdadelta(
+        var=var.handle,
+        accum=accum_grad.handle,
+        accum_update=accum_var.handle,
+        lr=coefficients['lr_t'],
+        rho=coefficients['rho'],
+        epsilon=coefficients['epsilon'],
+        grad=grad,
+        indices=indices,
+        use_locking=self._use_locking)
+
+  def get_config(self):
+    config = super(Adadelta, self).get_config()
+    config.update({
+        'learning_rate': self._serialize_hyperparameter('learning_rate'),
+        'decay': self._initial_decay,
+        'rho': self._serialize_hyperparameter('rho'),
+        'epsilon': self.epsilon,
+    })
+    return config

tensorflow/python/keras/optimizer_v2/adagrad.py

@@ -0,0 +1,170 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Adagrad optimizer implementation."""
+# pylint: disable=g-classes-have-attributes
+
+import numpy as np
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_conversion
+from tensorflow.python.keras import backend_config
+from tensorflow.python.keras.optimizer_v2 import optimizer_v2
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_training_ops
+from tensorflow.python.ops import init_ops
+
+
+class Adagrad(optimizer_v2.OptimizerV2):
+  r"""Optimizer that implements the Adagrad algorithm.
+
+  Adagrad is an optimizer with parameter-specific learning rates,
+  which are adapted relative to how frequently a parameter gets
+  updated during training. The more updates a parameter receives,
+  the smaller the updates.
+
+  Args:
+    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.
+      Note that `Adagrad` tends to benefit from higher initial learning rate
+      values compared to other optimizers.
+      To match the exact form in the original paper, use 1.0.
+    initial_accumulator_value: Floating point value.
+      Starting value for the accumulators (per-parameter momentum values).
+      Must be non-negative.
+    epsilon: Small floating point value used to maintain numerical stability.
+    name: Optional name prefix for the operations created when applying
+      gradients.  Defaults to `"Adagrad"`.
+    **kwargs: Keyword arguments. Allowed to be one of
+      `"clipnorm"` or `"clipvalue"`.
+      `"clipnorm"` (float) clips gradients by norm and represents
+      the maximum L2 norm of each weight variable;
+      `"clipvalue"` (float) clips gradient by value and represents the
+      maximum absolute value of each weight variable.
+
+  Reference:
+    - [Duchi et al., 2011](
+      http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf).
+  """
+
+  _HAS_AGGREGATE_GRAD = True
+
+  def __init__(self,
+               learning_rate=0.001,
+               initial_accumulator_value=0.1,
+               epsilon=1e-7,
+               name='Adagrad',
+               **kwargs):
+    if initial_accumulator_value < 0.0:
+      raise ValueError('initial_accumulator_value must be non-negative: %s' %
+                       initial_accumulator_value)
+    if epsilon is None:
+      epsilon = backend_config.epsilon()
+    super(Adagrad, self).__init__(name, **kwargs)
+    self._set_hyper('learning_rate', kwargs.get('lr', learning_rate))
+    self._set_hyper('decay', self._initial_decay)
+    self._initial_accumulator_value = initial_accumulator_value
+    self.epsilon = epsilon or backend_config.epsilon()
+
+  def _create_slots(self, var_list):
+    for var in var_list:
+      dtype = var.dtype.base_dtype
+      init = init_ops.constant_initializer(
+          self._initial_accumulator_value, dtype=dtype)
+      self.add_slot(var, 'accumulator', init)
+
+  def _prepare_local(self, var_device, var_dtype, apply_state):
+    super(Adagrad, self)._prepare_local(var_device, var_dtype, apply_state)
+    apply_state[(var_device, var_dtype)].update(
+        dict(
+            epsilon=tensor_conversion.convert_to_tensor_v2_with_dispatch(
+                self.epsilon, var_dtype
+            ),
+            neg_lr_t=-apply_state[(var_device, var_dtype)]['lr_t'],
+            zero=array_ops.zeros((), dtype=dtypes.int64),
+        )
+    )
+
+  def set_weights(self, weights):
+    params = self.weights
+    # Override set_weights for backward compatibility of Keras V1 optimizer
+    # since it does not include iteration at head of the weight list. Set
+    # iteration to 0.
+    if len(params) == len(weights) + 1:
+      weights = [np.array(0)] + weights
+    super(Adagrad, self).set_weights(weights)
+
+  @classmethod
+  def from_config(cls, config, custom_objects=None):
+    """Creates an optimizer from its config.
+
+    This method is the reverse of `get_config`,
+    capable of instantiating the same optimizer from the config
+    dictionary.
+
+    Args:
+        config: A Python dictionary, typically the output of get_config.
+        custom_objects: A Python dictionary mapping names to additional Python
+          objects used to create this optimizer, such as a function used for a
+          hyperparameter.
+
+    Returns:
+        An optimizer instance.
+    """
+    if 'initial_accumulator_value' not in config:
+      config['initial_accumulator_value'] = 0.1
+    if 'lr' in config:
+      config['learning_rate'] = config.pop('lr')
+    return cls(**config)
+
+  def _resource_apply_dense(self, grad, var, apply_state=None):
+    var_device, var_dtype = var.device, var.dtype.base_dtype
+    coefficients = ((apply_state or {}).get((var_device, var_dtype))
+                    or self._fallback_apply_state(var_device, var_dtype))
+
+    acc = self.get_slot(var, 'accumulator')
+    return gen_training_ops.ResourceApplyAdagradV2(
+        var=var.handle,
+        accum=acc.handle,
+        lr=coefficients['lr_t'],
+        epsilon=coefficients['epsilon'],
+        grad=grad,
+        use_locking=self._use_locking)
+
+  def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
+    var_device, var_dtype = var.device, var.dtype.base_dtype
+    coefficients = ((apply_state or {}).get((var_device, var_dtype))
+                    or self._fallback_apply_state(var_device, var_dtype))
+
+    acc = self.get_slot(var, 'accumulator')
+    return gen_training_ops.ResourceSparseApplyAdagradV2(
+        var=var.handle,
+        accum=acc.handle,
+        lr=coefficients['lr_t'],
+        epsilon=coefficients['epsilon'],
+        grad=grad,
+        indices=indices,
+        use_locking=self._use_locking)
+
+  def get_config(self):
+    config = super(Adagrad, self).get_config()
+    config.update({
+        'learning_rate': self._serialize_hyperparameter('learning_rate'),
+        'decay': self._initial_decay,
+        'initial_accumulator_value': self._initial_accumulator_value,
+        'epsilon': self.epsilon,
+    })
+    return config

tensorflow/python/keras/optimizer_v2/adam.py

@@ -0,0 +1,480 @@
+# Copyright 2020 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.
+# ==============================================================================
+"""Adam optimizer implementation."""
+# pylint: disable=g-classes-have-attributes
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_conversion
+from tensorflow.python.keras import backend_config
+from tensorflow.python.keras.optimizer_v2 import optimizer_v2
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import gen_training_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import state_ops
+
+
+class Adam(optimizer_v2.OptimizerV2):
+  r"""Optimizer that implements the Adam algorithm.
+
+  Adam optimization is a stochastic gradient descent method that is based on
+  adaptive estimation of first-order and second-order moments.
+
+  According to
+  [Kingma et al., 2014](http://arxiv.org/abs/1412.6980),
+  the method is "*computationally
+  efficient, has little memory requirement, invariant to diagonal rescaling of
+  gradients, and is well suited for problems that are large in terms of
+  data/parameters*".
+
+  Args:
+    learning_rate: A `Tensor`, floating point value, or a schedule that is a
+      `tf.keras.optimizers.schedules.LearningRateSchedule`, or a callable
+      that takes no arguments and returns the actual value to use, The
+      learning rate. Defaults to 0.001.
+    beta_1: A float value or a constant float tensor, or a callable
+      that takes no arguments and returns the actual value to use. The
+      exponential decay rate for the 1st moment estimates. Defaults to 0.9.
+    beta_2: A float value or a constant float tensor, or a callable
+      that takes no arguments and returns the actual value to use, The
+      exponential decay rate for the 2nd moment estimates. Defaults to 0.999.
+    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.
+    amsgrad: Boolean. Whether to apply AMSGrad variant of this algorithm from
+      the paper "On the Convergence of Adam and beyond". Defaults to `False`.
+    name: Optional name for the operations created when applying gradients.
+      Defaults to `"Adam"`.
+    **kwargs: Keyword arguments. Allowed to be one of
+      `"clipnorm"` or `"clipvalue"`.
+      `"clipnorm"` (float) clips gradients by norm; `"clipvalue"` (float) clips
+      gradients by value.
+
+  Usage:
+
+  >>> opt = tf.keras.optimizers.Adam(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()
+  >>> # The first step is `-learning_rate*sign(grad)`
+  >>> var1.numpy()
+  9.9
+
+  Reference:
+    - [Kingma et al., 2014](http://arxiv.org/abs/1412.6980)
+    - [Reddi et al., 2018](
+        https://openreview.net/pdf?id=ryQu7f-RZ) for `amsgrad`.
+
+  Notes:
+
+  The default value of 1e-7 for epsilon might not be a good default in
+  general. For example, when training an Inception network on ImageNet a
+  current good choice is 1.0 or 0.1. Note that since Adam uses the
+  formulation just before Section 2.1 of the Kingma and Ba paper rather than
+  the formulation in Algorithm 1, the "epsilon" referred to here is "epsilon
+  hat" in the paper.
+
+  The sparse implementation of this algorithm (used when the gradient is an
+  IndexedSlices object, typically because of `tf.gather` or an embedding
+  lookup in the forward pass) does apply momentum to variable slices even if
+  they were not used in the forward pass (meaning they have a gradient equal
+  to zero). Momentum decay (beta1) is also applied to the entire momentum
+  accumulator. This means that the sparse behavior is equivalent to the dense
+  behavior (in contrast to some momentum implementations which ignore momentum
+  unless a variable slice was actually used).
+  """
+
+  _HAS_AGGREGATE_GRAD = True
+
+  def __init__(self,
+               learning_rate=0.001,
+               beta_1=0.9,
+               beta_2=0.999,
+               epsilon=1e-7,
+               amsgrad=False,
+               name='Adam',
+               **kwargs):
+    super(Adam, self).__init__(name, **kwargs)
+    self._set_hyper('learning_rate', kwargs.get('lr', learning_rate))
+    self._set_hyper('decay', self._initial_decay)
+    self._set_hyper('beta_1', beta_1)
+    self._set_hyper('beta_2', beta_2)
+    self.epsilon = epsilon or backend_config.epsilon()
+    self.amsgrad = amsgrad
+
+  def _create_slots(self, var_list):
+    # Create slots for the first and second moments.
+    # Separate for-loops to respect the ordering of slot variables from v1.
+    for var in var_list:
+      self.add_slot(var, 'm')
+    for var in var_list:
+      self.add_slot(var, 'v')
+    if self.amsgrad:
+      for var in var_list:
+        self.add_slot(var, 'vhat')
+
+  def _prepare_local(self, var_device, var_dtype, apply_state):
+    super(Adam, self)._prepare_local(var_device, var_dtype, apply_state)
+
+    local_step = math_ops.cast(self.iterations + 1, var_dtype)
+    beta_1_t = array_ops.identity(self._get_hyper('beta_1', var_dtype))
+    beta_2_t = array_ops.identity(self._get_hyper('beta_2', var_dtype))
+    beta_1_power = math_ops.pow(beta_1_t, local_step)
+    beta_2_power = math_ops.pow(beta_2_t, local_step)
+    lr = (apply_state[(var_device, var_dtype)]['lr_t'] *
+          (math_ops.sqrt(1 - beta_2_power) / (1 - beta_1_power)))
+    apply_state[(var_device, var_dtype)].update(
+        dict(
+            lr=lr,
+            epsilon=tensor_conversion.convert_to_tensor_v2_with_dispatch(
+                self.epsilon, var_dtype
+            ),
+            beta_1_t=beta_1_t,
+            beta_1_power=beta_1_power,
+            one_minus_beta_1_t=1 - beta_1_t,
+            beta_2_t=beta_2_t,
+            beta_2_power=beta_2_power,
+            one_minus_beta_2_t=1 - beta_2_t,
+        )
+    )
+
+  def set_weights(self, weights):
+    params = self.weights
+    # If the weights are generated by Keras V1 optimizer, it includes vhats
+    # even without amsgrad, i.e, V1 optimizer has 3x + 1 variables, while V2
+    # optimizer has 2x + 1 variables. Filter vhats out for compatibility.
+    num_vars = int((len(params) - 1) / 2)
+    if len(weights) == 3 * num_vars + 1:
+      weights = weights[:len(params)]
+    super(Adam, self).set_weights(weights)
+
+  def _resource_apply_dense(self, grad, var, apply_state=None):
+    var_device, var_dtype = var.device, var.dtype.base_dtype
+    coefficients = ((apply_state or {}).get((var_device, var_dtype))
+                    or self._fallback_apply_state(var_device, var_dtype))
+
+    m = self.get_slot(var, 'm')
+    v = self.get_slot(var, 'v')
+
+    if not self.amsgrad:
+      return gen_training_ops.ResourceApplyAdam(
+          var=var.handle,
+          m=m.handle,
+          v=v.handle,
+          beta1_power=coefficients['beta_1_power'],
+          beta2_power=coefficients['beta_2_power'],
+          lr=coefficients['lr_t'],
+          beta1=coefficients['beta_1_t'],
+          beta2=coefficients['beta_2_t'],
+          epsilon=coefficients['epsilon'],
+          grad=grad,
+          use_locking=self._use_locking)
+    else:
+      vhat = self.get_slot(var, 'vhat')
+      return gen_training_ops.ResourceApplyAdamWithAmsgrad(
+          var=var.handle,
+          m=m.handle,
+          v=v.handle,
+          vhat=vhat.handle,
+          beta1_power=coefficients['beta_1_power'],
+          beta2_power=coefficients['beta_2_power'],
+          lr=coefficients['lr_t'],
+          beta1=coefficients['beta_1_t'],
+          beta2=coefficients['beta_2_t'],
+          epsilon=coefficients['epsilon'],
+          grad=grad,
+          use_locking=self._use_locking)
+
+  def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
+    var_device, var_dtype = var.device, var.dtype.base_dtype
+    coefficients = ((apply_state or {}).get((var_device, var_dtype))
+                    or self._fallback_apply_state(var_device, var_dtype))
+
+    # m_t = beta1 * m + (1 - beta1) * g_t
+    m = self.get_slot(var, 'm')
+    m_scaled_g_values = grad * coefficients['one_minus_beta_1_t']
+    m_t = state_ops.assign(m, m * coefficients['beta_1_t'],
+                           use_locking=self._use_locking)
+    with ops.control_dependencies([m_t]):
+      m_t = self._resource_scatter_add(m, indices, m_scaled_g_values)
+
+    # v_t = beta2 * v + (1 - beta2) * (g_t * g_t)
+    v = self.get_slot(var, 'v')
+    v_scaled_g_values = (grad * grad) * coefficients['one_minus_beta_2_t']
+    v_t = state_ops.assign(v, v * coefficients['beta_2_t'],
+                           use_locking=self._use_locking)
+    with ops.control_dependencies([v_t]):
+      v_t = self._resource_scatter_add(v, indices, v_scaled_g_values)
+
+    if not self.amsgrad:
+      v_sqrt = math_ops.sqrt(v_t)
+      var_update = state_ops.assign_sub(
+          var, coefficients['lr'] * m_t / (v_sqrt + coefficients['epsilon']),
+          use_locking=self._use_locking)
+      return control_flow_ops.group(*[var_update, m_t, v_t])
+    else:
+      v_hat = self.get_slot(var, 'vhat')
+      v_hat_t = math_ops.maximum(v_hat, v_t)
+      with ops.control_dependencies([v_hat_t]):
+        v_hat_t = state_ops.assign(
+            v_hat, v_hat_t, use_locking=self._use_locking)
+      v_hat_sqrt = math_ops.sqrt(v_hat_t)
+      var_update = state_ops.assign_sub(
+          var,
+          coefficients['lr'] * m_t / (v_hat_sqrt + coefficients['epsilon']),
+          use_locking=self._use_locking)
+      return control_flow_ops.group(*[var_update, m_t, v_t, v_hat_t])
+
+  def get_config(self):
+    config = super(Adam, self).get_config()
+    config.update({
+        'learning_rate': self._serialize_hyperparameter('learning_rate'),
+        'decay': self._initial_decay,
+        'beta_1': self._serialize_hyperparameter('beta_1'),
+        'beta_2': self._serialize_hyperparameter('beta_2'),
+        'epsilon': self.epsilon,
+        'amsgrad': self.amsgrad,
+    })
+    return config
+
+
+class NonFusedAdam(optimizer_v2.OptimizerV2):
+  r"""Optimizer that implements the Adam algorithm without fused kernels.
+
+  Adam optimization is a stochastic gradient descent method that is based on
+  adaptive estimation of first-order and second-order moments.
+  According to the paper
+  [Adam: A Method for Stochastic Optimization. Kingma et al.,
+  2014](http://arxiv.org/abs/1412.6980), the method is "*computationally
+  efficient, has little memory requirement, invariant to diagonal rescaling of
+  gradients, and is well suited for problems that are large in terms of
+  data/parameters*".
+
+  For AMSGrad see [On The Convergence Of Adam And Beyond.
+  Reddi et al., 5-8](https://openreview.net/pdf?id=ryQu7f-RZ).
+
+  **If amsgrad = False**:
+
+  initialize $m_0$ as 1st moment vector
+  initialize $v_0$ as 2nd moment vector
+
+  The update rule for $\theta$ with gradient $g$ uses an optimization
+  described at the end of section 2 of the paper:
+
+  $$lr_t = \mathrm{learning\_rate} *
+    \sqrt{1 - \beta_2^t} / (1 - \beta_1^t)$$
+  $$m_t = \beta_1 * m_{t-1} + (1 - \beta_1) * g$$
+  $$v_t = \beta_2 * v_{t-1} + (1 - \beta_2) * g^2$$
+  $$\theta_t = \theta_{t-1} - lr_t * m_t / (\sqrt{v_t} + \epsilon)$$
+
+  **If amsgrad = True**:
+
+  initialize $m_0$ as 1st moment vector
+  initialize $v_0$ as 2nd moment vector
+  initialize $\hat{v}_0$ as 2nd moment vector
+
+  The update rule for $\theta$ with gradient $g$ uses an optimization
+  described at the end of section 2 of the paper:
+
+  $$lr_t = \mathrm{learning\_rate} *
+    \sqrt{1 - \beta_2^t} / (1 - \beta_1^t)$$
+
+  $$m_t = \beta_1 * m_{t-1} + (1 - \beta_1) * g$$
+  $$v_t = \beta_2 * v_{t-1} + (1 - \beta_2) * g^2$$
+  $$\hat{v}_t = \max(\hat{v}_{t-1}, v_t)$$
+  $$\theta_t = \theta_{t-1} - lr_t * m_t / (\sqrt{\hat{v}_t} + \epsilon)$$
+
+  The default value of 1e-7 for epsilon might not be a good default in
+  general. For example, when training an Inception network on ImageNet a
+  current good choice is 1.0 or 0.1. Note that since Adam uses the
+  formulation just before Section 2.1 of the Kingma and Ba paper rather than
+  the formulation in Algorithm 1, the "epsilon" referred to here is "epsilon
+  hat" in the paper.
+
+  The sparse implementation of this algorithm (used when the gradient is an
+  IndexedSlices object, typically because of `tf.gather` or an embedding
+  lookup in the forward pass) does apply momentum to variable slices even if
+  they were not used in the forward pass (meaning they have a gradient equal
+  to zero). Momentum decay (beta1) is also applied to the entire momentum
+  accumulator. This means that the sparse behavior is equivalent to the dense
+  behavior (in contrast to some momentum implementations which ignore momentum
+  unless a variable slice was actually used).
+
+  Usage:
+
+  >>> opt = tf.keras.optimizers.Adam(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()
+  >>> # The first step is `-learning_rate*sign(grad)`
+  >>> var1.numpy()
+  9.9
+  """
+
+  _HAS_AGGREGATE_GRAD = True
+
+  def __init__(self,
+               learning_rate=0.001,
+               beta_1=0.9,
+               beta_2=0.999,
+               epsilon=1e-7,
+               amsgrad=False,
+               name='Adam',
+               **kwargs):
+    """Construct a new Adam optimizer.
+
+    Args:
+      learning_rate: A `Tensor`, floating point value, or a schedule that is a
+        `tf.keras.optimizers.schedules.LearningRateSchedule`, or a callable that
+        takes no arguments and returns the actual value to use, The learning
+        rate. Defaults to 0.001.
+      beta_1: A float value or a constant float tensor, or a callable that takes
+        no arguments and returns the actual value to use. The exponential decay
+        rate for the 1st moment estimates. Defaults to 0.9.
+      beta_2: A float value or a constant float tensor, or a callable that takes
+        no arguments and returns the actual value to use, The exponential decay
+        rate for the 2nd moment estimates. Defaults to 0.999.
+      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.
+      amsgrad: Boolean. Whether to apply AMSGrad variant of this algorithm from
+        the paper "On the Convergence of Adam and beyond". Defaults to `False`.
+      name: Optional name for the operations created when applying gradients.
+        Defaults to "Adam".
+      **kwargs: keyword arguments. Allowed to be {`clipnorm`, `clipvalue`, `lr`,
+        `decay`}. `clipnorm` is clip gradients by norm; `clipvalue` is clip
+        gradients by value, `decay` is included for backward compatibility to
+        allow time inverse decay of learning rate. `lr` is included for backward
+        compatibility, recommended to use `learning_rate` instead.
+    """
+
+    super(NonFusedAdam, self).__init__(name, **kwargs)
+    self._set_hyper('learning_rate', kwargs.get('lr', learning_rate))
+    self._set_hyper('decay', self._initial_decay)
+    self._set_hyper('beta_1', beta_1)
+    self._set_hyper('beta_2', beta_2)
+    self.epsilon = epsilon or backend_config.epsilon()
+    self.amsgrad = amsgrad
+
+  def _create_slots(self, var_list):
+    # Create slots for the first and second moments.
+    # Separate for-loops to respect the ordering of slot variables from v1.
+    for var in var_list:
+      self.add_slot(var, 'm')
+    for var in var_list:
+      self.add_slot(var, 'v')
+    if self.amsgrad:
+      for var in var_list:
+        self.add_slot(var, 'vhat')
+
+  def _prepare_local(self, var_device, var_dtype, apply_state):
+    super(NonFusedAdam, self)._prepare_local(var_device, var_dtype, apply_state)
+
+    local_step = math_ops.cast(self.iterations + 1, var_dtype)
+    beta_1_t = array_ops.identity(self._get_hyper('beta_1', var_dtype))
+    beta_2_t = array_ops.identity(self._get_hyper('beta_2', var_dtype))
+    beta_1_power = math_ops.pow(beta_1_t, local_step)
+    beta_2_power = math_ops.pow(beta_2_t, local_step)
+    lr = (
+        apply_state[(var_device, var_dtype)]['lr_t'] *
+        (math_ops.sqrt(1 - beta_2_power) / (1 - beta_1_power)))
+    apply_state[(var_device, var_dtype)].update(
+        dict(
+            lr=lr,
+            epsilon=tensor_conversion.convert_to_tensor_v2_with_dispatch(
+                self.epsilon, var_dtype
+            ),
+            beta_1_t=beta_1_t,
+            beta_1_power=beta_1_power,
+            one_minus_beta_1_t=1 - beta_1_t,
+            beta_2_t=beta_2_t,
+            beta_2_power=beta_2_power,
+            one_minus_beta_2_t=1 - beta_2_t,
+        )
+    )
+
+  def set_weights(self, weights):
+    params = self.weights
+    # If the weights are generated by Keras V1 optimizer, it includes vhats
+    # even without amsgrad, i.e, V1 optimizer has 3x + 1 variables, while V2
+    # optimizer has 2x + 1 variables. Filter vhats out for compatibility.
+    num_vars = int((len(params) - 1) / 2)
+    if len(weights) == 3 * num_vars + 1:
+      weights = weights[:len(params)]
+    super(NonFusedAdam, self).set_weights(weights)
+
+  @def_function.function(jit_compile=True)
+  def _resource_apply_dense(self, grad, var, apply_state=None):
+    var_device, var_dtype = var.device, var.dtype.base_dtype
+    coefficients = ((apply_state or {}).get((var_device, var_dtype)) or
+                    self._fallback_apply_state(var_device, var_dtype))
+
+    m = self.get_slot(var, 'm')
+    v = self.get_slot(var, 'v')
+
+    alpha = (
+        coefficients['lr_t'] * math_ops.sqrt(1 - coefficients['beta_2_power']) /
+        (1 - coefficients['beta_1_power']))
+    m.assign_add((grad - m) * (1 - coefficients['beta_1_t']))
+    v.assign_add((math_ops.square(grad) - v) * (1 - coefficients['beta_2_t']))
+    if self.amsgrad:
+      vhat = self.get_slot(var, 'vhat')
+      vhat.assign(math_ops.maximum(vhat, v))
+      v = vhat
+    var.assign_sub(
+        (m * alpha) / (math_ops.sqrt(v) - coefficients['epsilon']))
+
+  @def_function.function(jit_compile=True)
+  def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
+    var_device, var_dtype = var.device, var.dtype.base_dtype
+    coefficients = ((apply_state or {}).get((var_device, var_dtype)) or
+                    self._fallback_apply_state(var_device, var_dtype))
+
+    # m_t = beta1 * m + (1 - beta1) * g_t
+    m = self.get_slot(var, 'm')
+    m_scaled_g_values = grad * coefficients['one_minus_beta_1_t']
+    m.assign(m * coefficients['beta_1_t'])
+    m.scatter_add(indexed_slices.IndexedSlices(m_scaled_g_values, indices))
+
+    # v_t = beta2 * v + (1 - beta2) * (g_t * g_t)
+    v = self.get_slot(var, 'v')
+    v_scaled_g_values = (grad * grad) * coefficients['one_minus_beta_2_t']
+    v.assign(v * coefficients['beta_2_t'])
+    v.scatter_add(indexed_slices.IndexedSlices(v_scaled_g_values, indices))
+
+    if not self.amsgrad:
+      var.assign_sub(coefficients['lr'] * m /
+                     (math_ops.sqrt(v) + coefficients['epsilon']))
+    else:
+      v_hat = self.get_slot(var, 'vhat')
+      v_hat.assign(math_ops.maximum(v_hat, v))
+      var.assign_sub(coefficients['lr'] * m /
+                     (math_ops.sqrt(v_hat) + coefficients['epsilon']))
+
+  def get_config(self):
+    config = super(NonFusedAdam, self).get_config()
+    config.update({
+        'learning_rate': self._serialize_hyperparameter('learning_rate'),
+        'decay': self._initial_decay,
+        'beta_1': self._serialize_hyperparameter('beta_1'),
+        'beta_2': self._serialize_hyperparameter('beta_2'),
+        'epsilon': self.epsilon,
+        'amsgrad': self.amsgrad,
+    })
+    return config

tensorflow/python/keras/optimizer_v2/adamax.py

@@ -0,0 +1,187 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Adamax optimizer implementation."""
+# pylint: disable=g-classes-have-attributes
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_conversion
+from tensorflow.python.keras import backend_config
+from tensorflow.python.keras.optimizer_v2 import optimizer_v2
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import gen_training_ops
+from tensorflow.python.ops import math_ops
+
+
+class Adamax(optimizer_v2.OptimizerV2):
+  """Optimizer that implements the Adamax algorithm.
+
+  It is a variant of Adam based on the infinity norm.
+  Default parameters follow those provided in the paper.
+  Adamax is sometimes superior to adam, specially in models with embeddings.
+
+  Initialization:
+
+  ```python
+  m = 0  # Initialize initial 1st moment vector
+  v = 0  # Initialize the exponentially weighted infinity norm
+  t = 0  # Initialize timestep
+  ```
+
+  The update rule for parameter `w` with gradient `g` is
+  described at the end of section 7.1 of the paper:
+
+  ```python
+  t += 1
+  m = beta1 * m + (1 - beta) * g
+  v = max(beta2 * v, abs(g))
+  current_lr = learning_rate / (1 - beta1 ** t)
+  w = w - current_lr * m / (v + epsilon)
+  ```
+
+  Similarly to `Adam`, the epsilon is added for numerical stability
+  (especially to get rid of division by zero when `v_t == 0`).
+
+  In contrast to `Adam`, the sparse implementation of this algorithm
+  (used when the gradient is an IndexedSlices object, typically because of
+  `tf.gather` or an embedding lookup in the forward pass) only updates
+  variable slices and corresponding `m_t`, `v_t` terms when that part of
+  the variable was used in the forward pass. This means that the sparse
+  behavior is contrast to the dense behavior (similar to some momentum
+  implementations which ignore momentum unless a variable slice was actually
+  used).
+
+  Args:
+    learning_rate: A `Tensor`, floating point value, or a schedule that is a
+      `tf.keras.optimizers.schedules.LearningRateSchedule`. The learning rate.
+    beta_1: A float value or a constant float tensor. The exponential decay
+      rate for the 1st moment estimates.
+    beta_2: A float value or a constant float tensor. The exponential decay
+      rate for the exponentially weighted infinity norm.
+    epsilon: A small constant for numerical stability.
+    name: Optional name for the operations created when applying gradients.
+      Defaults to `"Adamax"`.
+    **kwargs: Keyword arguments. Allowed to be one of
+      `"clipnorm"` or `"clipvalue"`.
+      `"clipnorm"` (float) clips gradients by norm; `"clipvalue"` (float) clips
+      gradients by value.
+
+  Reference:
+    - [Kingma et al., 2014](http://arxiv.org/abs/1412.6980)
+  """
+
+  _HAS_AGGREGATE_GRAD = True
+
+  def __init__(self,
+               learning_rate=0.001,
+               beta_1=0.9,
+               beta_2=0.999,
+               epsilon=1e-7,
+               name='Adamax',
+               **kwargs):
+    super(Adamax, self).__init__(name, **kwargs)
+    self._set_hyper('learning_rate', kwargs.get('lr', learning_rate))
+    self._set_hyper('decay', self._initial_decay)
+    self._set_hyper('beta_1', beta_1)
+    self._set_hyper('beta_2', beta_2)
+    self.epsilon = epsilon or backend_config.epsilon()
+
+  def _create_slots(self, var_list):
+    # Separate for-loops to respect the ordering of slot variables from v1.
+    for var in var_list:
+      self.add_slot(var, 'm')  # Create slots for the first moments.
+    for var in var_list:
+      self.add_slot(var, 'v')  # Create slots for the second moments.
+
+  def _prepare_local(self, var_device, var_dtype, apply_state):
+    super(Adamax, self)._prepare_local(var_device, var_dtype, apply_state)
+
+    local_step = math_ops.cast(self.iterations + 1, var_dtype)
+    beta_1_t = array_ops.identity(self._get_hyper('beta_1', var_dtype))
+    beta_2_t = array_ops.identity(self._get_hyper('beta_2', var_dtype))
+    beta_1_power = math_ops.pow(beta_1_t, local_step)
+    lr_t = apply_state[(var_device, var_dtype)]['lr_t']
+
+    apply_state[(var_device, var_dtype)].update(
+        dict(
+            neg_scaled_lr=-lr_t / (1 - beta_1_power),
+            epsilon=tensor_conversion.convert_to_tensor_v2_with_dispatch(
+                self.epsilon, var_dtype
+            ),
+            beta_1_t=beta_1_t,
+            beta_1_power=beta_1_power,
+            one_minus_beta_1_t=1 - beta_1_t,
+            beta_2_t=beta_2_t,
+            zero=array_ops.zeros((), dtype=dtypes.int64),
+        )
+    )
+
+  def _resource_apply_dense(self, grad, var, apply_state=None):
+    var_device, var_dtype = var.device, var.dtype.base_dtype
+    coefficients = ((apply_state or {}).get((var_device, var_dtype))
+                    or self._fallback_apply_state(var_device, var_dtype))
+
+    m = self.get_slot(var, 'm')
+    v = self.get_slot(var, 'v')
+    return gen_training_ops.ResourceApplyAdaMax(
+        var=var.handle,
+        m=m.handle,
+        v=v.handle,
+        beta1_power=coefficients['beta_1_power'],
+        lr=coefficients['lr_t'],
+        beta1=coefficients['beta_1_t'],
+        beta2=coefficients['beta_2_t'],
+        epsilon=coefficients['epsilon'],
+        grad=grad,
+        use_locking=self._use_locking)
+
+  def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
+    var_device, var_dtype = var.device, var.dtype.base_dtype
+    coefficients = ((apply_state or {}).get((var_device, var_dtype))
+                    or self._fallback_apply_state(var_device, var_dtype))
+
+    # m_t = beta1 * m + (1 - beta1) * g_t
+    m = self.get_slot(var, 'm')
+    m_slice = array_ops.gather(m, indices, axis=coefficients['zero'])
+    m_t_slice = (m_slice * coefficients['beta_1_t'] +
+                 grad * coefficients['one_minus_beta_1_t'])
+    with ops.control_dependencies([m_t_slice]):
+      m_t = self._resource_scatter_update(m, indices, m_t_slice)
+
+    # u_t = max(beta2 * u, abs(g_t))
+    v = self.get_slot(var, 'v')
+    v_slice = array_ops.gather(v, indices, axis=coefficients['zero'])
+    v_t_slice = math_ops.maximum(v_slice * coefficients['beta_2_t'],
+                                 math_ops.abs(grad))
+    with ops.control_dependencies([v_t_slice]):
+      v_t = self._resource_scatter_update(v, indices, v_t_slice)
+    # theta_t = theta - lr / (1 - beta1^t) * m_t / u_t
+    var_slice = coefficients['neg_scaled_lr'] * (
+        m_t_slice / (v_t_slice + coefficients['epsilon']))
+    with ops.control_dependencies([var_slice]):
+      var_update = self._resource_scatter_add(var, indices, var_slice)
+    return control_flow_ops.group(*[var_update, m_t, v_t])
+
+  def get_config(self):
+    config = super(Adamax, self).get_config()
+    config.update({
+        'learning_rate': self._serialize_hyperparameter('learning_rate'),
+        'decay': self._initial_decay,
+        'beta_1': self._serialize_hyperparameter('beta_1'),
+        'beta_2': self._serialize_hyperparameter('beta_2'),
+        'epsilon': self.epsilon,
+    })
+    return config

tensorflow/python/keras/optimizer_v2/ftrl.py

@@ -0,0 +1,263 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Ftrl-proximal optimizer implementation."""
+# pylint: disable=g-classes-have-attributes
+
+from tensorflow.python.keras.optimizer_v2 import optimizer_v2
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_training_ops
+from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import math_ops
+
+
+class Ftrl(optimizer_v2.OptimizerV2):
+  r"""Optimizer that implements the FTRL algorithm.
+
+  "Follow The Regularized Leader" (FTRL) is an optimization algorithm developed
+  at Google for click-through rate prediction in the early 2010s. It is most
+  suitable for shallow models with large and sparse feature spaces.
+  The algorithm is described by
+  [McMahan et al., 2013](https://research.google.com/pubs/archive/41159.pdf).
+  The Keras version has support for both online L2 regularization
+  (the L2 regularization described in the paper
+  above) and shrinkage-type L2 regularization
+  (which is the addition of an L2 penalty to the loss function).
+
+  Initialization:
+
+  ```python
+  n = 0
+  sigma = 0
+  z = 0
+  ```
+
+  Update rule for one variable `w`:
+
+  ```python
+  prev_n = n
+  n = n + g ** 2
+  sigma = (sqrt(n) - sqrt(prev_n)) / lr
+  z = z + g - sigma * w
+  if abs(z) < lambda_1:
+    w = 0
+  else:
+    w = (sgn(z) * lambda_1 - z) / ((beta + sqrt(n)) / alpha + lambda_2)
+  ```
+
+  Notation:
+
+  - `lr` is the learning rate
+  - `g` is the gradient for the variable
+  - `lambda_1` is the L1 regularization strength
+  - `lambda_2` is the L2 regularization strength
+
+  Check the documentation for the `l2_shrinkage_regularization_strength`
+  parameter for more details when shrinkage is enabled, in which case gradient
+  is replaced with a gradient with shrinkage.
+
+  Args:
+    learning_rate: A `Tensor`, floating point value, or a schedule that is a
+      `tf.keras.optimizers.schedules.LearningRateSchedule`. The learning rate.
+    learning_rate_power: A float value, must be less or equal to zero.
+      Controls how the learning rate decreases during training. Use zero for
+      a fixed learning rate.
+    initial_accumulator_value: The starting value for accumulators.
+      Only zero or positive values are allowed.
+    l1_regularization_strength: A float value, must be greater than or
+      equal to zero. Defaults to 0.0.
+    l2_regularization_strength: A float value, must be greater than or
+      equal to zero. Defaults to 0.0.
+    name: Optional name prefix for the operations created when applying
+      gradients.  Defaults to `"Ftrl"`.
+    l2_shrinkage_regularization_strength: A float value, must be greater than
+      or equal to zero. This differs from L2 above in that the L2 above is a
+      stabilization penalty, whereas this L2 shrinkage is a magnitude penalty.
+      When input is sparse shrinkage will only happen on the active weights.
+    beta: A float value, representing the beta value from the paper.
+      Defaults to 0.0.
+    **kwargs: Keyword arguments. Allowed to be one of
+      `"clipnorm"` or `"clipvalue"`.
+      `"clipnorm"` (float) clips gradients by norm; `"clipvalue"` (float) clips
+      gradients by value.
+
+  Reference:
+    - [McMahan et al., 2013](
+      https://research.google.com/pubs/archive/41159.pdf)
+  """
+
+  def __init__(self,
+               learning_rate=0.001,
+               learning_rate_power=-0.5,
+               initial_accumulator_value=0.1,
+               l1_regularization_strength=0.0,
+               l2_regularization_strength=0.0,
+               name='Ftrl',
+               l2_shrinkage_regularization_strength=0.0,
+               beta=0.0,
+               **kwargs):
+    super(Ftrl, self).__init__(name, **kwargs)
+
+    if initial_accumulator_value < 0.0:
+      raise ValueError(
+          'initial_accumulator_value %f needs to be positive or zero' %
+          initial_accumulator_value)
+    if learning_rate_power > 0.0:
+      raise ValueError('learning_rate_power %f needs to be negative or zero' %
+                       learning_rate_power)
+    if l1_regularization_strength < 0.0:
+      raise ValueError(
+          'l1_regularization_strength %f needs to be positive or zero' %
+          l1_regularization_strength)
+    if l2_regularization_strength < 0.0:
+      raise ValueError(
+          'l2_regularization_strength %f needs to be positive or zero' %
+          l2_regularization_strength)
+    if l2_shrinkage_regularization_strength < 0.0:
+      raise ValueError(
+          'l2_shrinkage_regularization_strength %f needs to be positive'
+          ' or zero' % l2_shrinkage_regularization_strength)
+
+    self._set_hyper('learning_rate', learning_rate)
+    self._set_hyper('decay', self._initial_decay)
+    self._set_hyper('learning_rate_power', learning_rate_power)
+    self._set_hyper('l1_regularization_strength', l1_regularization_strength)
+    self._set_hyper('l2_regularization_strength', l2_regularization_strength)
+    self._set_hyper('beta', beta)
+    self._initial_accumulator_value = initial_accumulator_value
+    self._l2_shrinkage_regularization_strength = (
+        l2_shrinkage_regularization_strength)
+
+  def _create_slots(self, var_list):
+    # Create the "accum" and "linear" slots.
+    for var in var_list:
+      dtype = var.dtype.base_dtype
+      init = init_ops.constant_initializer(
+          self._initial_accumulator_value, dtype=dtype)
+      self.add_slot(var, 'accumulator', init)
+      self.add_slot(var, 'linear')
+
+  def _prepare_local(self, var_device, var_dtype, apply_state):
+    super(Ftrl, self)._prepare_local(var_device, var_dtype, apply_state)
+    apply_state[(var_device, var_dtype)].update(
+        dict(
+            learning_rate_power=array_ops.identity(
+                self._get_hyper('learning_rate_power', var_dtype)),
+            l1_regularization_strength=array_ops.identity(
+                self._get_hyper('l1_regularization_strength', var_dtype)),
+            l2_regularization_strength=array_ops.identity(
+                self._get_hyper('l2_regularization_strength', var_dtype)),
+            beta=array_ops.identity(self._get_hyper('beta', var_dtype)),
+            l2_shrinkage_regularization_strength=math_ops.cast(
+                self._l2_shrinkage_regularization_strength, var_dtype)))
+
+  def _resource_apply_dense(self, grad, var, apply_state=None):
+    var_device, var_dtype = var.device, var.dtype.base_dtype
+    coefficients = ((apply_state or {}).get((var_device, var_dtype))
+                    or self._fallback_apply_state(var_device, var_dtype))
+
+    # Adjust L2 regularization strength to include beta to avoid the underlying
+    # TensorFlow ops needing to include it.
+    adjusted_l2_regularization_strength = (
+        coefficients['l2_regularization_strength'] + coefficients['beta'] /
+        (2. * coefficients['lr_t']))
+
+    accum = self.get_slot(var, 'accumulator')
+    linear = self.get_slot(var, 'linear')
+
+    if self._l2_shrinkage_regularization_strength <= 0.0:
+      return gen_training_ops.ResourceApplyFtrl(
+          var=var.handle,
+          accum=accum.handle,
+          linear=linear.handle,
+          grad=grad,
+          lr=coefficients['lr_t'],
+          l1=coefficients['l1_regularization_strength'],
+          l2=adjusted_l2_regularization_strength,
+          lr_power=coefficients['learning_rate_power'],
+          use_locking=self._use_locking)
+    else:
+      return gen_training_ops.ResourceApplyFtrlV2(
+          var=var.handle,
+          accum=accum.handle,
+          linear=linear.handle,
+          grad=grad,
+          lr=coefficients['lr_t'],
+          l1=coefficients['l1_regularization_strength'],
+          l2=adjusted_l2_regularization_strength,
+          l2_shrinkage=coefficients['l2_shrinkage_regularization_strength'],
+          lr_power=coefficients['learning_rate_power'],
+          use_locking=self._use_locking)
+
+  def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
+    var_device, var_dtype = var.device, var.dtype.base_dtype
+    coefficients = ((apply_state or {}).get((var_device, var_dtype))
+                    or self._fallback_apply_state(var_device, var_dtype))
+
+    # Adjust L2 regularization strength to include beta to avoid the underlying
+    # TensorFlow ops needing to include it.
+    adjusted_l2_regularization_strength = (
+        coefficients['l2_regularization_strength'] + coefficients['beta'] /
+        (2. * coefficients['lr_t']))
+
+    accum = self.get_slot(var, 'accumulator')
+    linear = self.get_slot(var, 'linear')
+
+    if self._l2_shrinkage_regularization_strength <= 0.0:
+      return gen_training_ops.ResourceSparseApplyFtrl(
+          var=var.handle,
+          accum=accum.handle,
+          linear=linear.handle,
+          grad=grad,
+          indices=indices,
+          lr=coefficients['lr_t'],
+          l1=coefficients['l1_regularization_strength'],
+          l2=adjusted_l2_regularization_strength,
+          lr_power=coefficients['learning_rate_power'],
+          use_locking=self._use_locking)
+    else:
+      return gen_training_ops.ResourceSparseApplyFtrlV2(
+          var=var.handle,
+          accum=accum.handle,
+          linear=linear.handle,
+          grad=grad,
+          indices=indices,
+          lr=coefficients['lr_t'],
+          l1=coefficients['l1_regularization_strength'],
+          l2=adjusted_l2_regularization_strength,
+          l2_shrinkage=coefficients['l2_shrinkage_regularization_strength'],
+          lr_power=coefficients['learning_rate_power'],
+          use_locking=self._use_locking)
+
+  def get_config(self):
+    config = super(Ftrl, self).get_config()
+    config.update({
+        'learning_rate':
+            self._serialize_hyperparameter('learning_rate'),
+        'decay':
+            self._initial_decay,
+        'initial_accumulator_value':
+            self._initial_accumulator_value,
+        'learning_rate_power':
+            self._serialize_hyperparameter('learning_rate_power'),
+        'l1_regularization_strength':
+            self._serialize_hyperparameter('l1_regularization_strength'),
+        'l2_regularization_strength':
+            self._serialize_hyperparameter('l2_regularization_strength'),
+        'beta':
+            self._serialize_hyperparameter('beta'),
+        'l2_shrinkage_regularization_strength':
+            self._l2_shrinkage_regularization_strength,
+    })
+    return config

tensorflow/python/keras/optimizer_v2/gradient_descent.py

@@ -0,0 +1,190 @@
+# Copyright 2020 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.
+# ==============================================================================
+"""SGD optimizer implementation."""
+# pylint: disable=g-classes-have-attributes
+
+from tensorflow.python.framework import tensor
+from tensorflow.python.keras.optimizer_v2 import optimizer_v2
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_resource_variable_ops
+from tensorflow.python.ops import gen_training_ops
+
+
+class SGD(optimizer_v2.OptimizerV2):
+  r"""Gradient descent (with momentum) optimizer.
+
+  Update rule for parameter `w` with gradient `g` when `momentum` is 0:
+
+  ```python
+  w = w - learning_rate * g
+  ```
+
+  Update rule when `momentum` is larger than 0:
+
+  ```python
+  velocity = momentum * velocity - learning_rate * g
+  w = w + velocity
+  ```
+
+  When `nesterov=True`, this rule becomes:
+
+  ```python
+  velocity = momentum * velocity - learning_rate * g
+  w = w + momentum * velocity - learning_rate * g
+  ```
+
+  Args:
+    learning_rate: A `Tensor`, floating point value, or a schedule that is a
+      `tf.keras.optimizers.schedules.LearningRateSchedule`, or a callable
+      that takes no arguments and returns the actual value to use. The
+      learning rate. Defaults to 0.01.
+    momentum: float hyperparameter >= 0 that accelerates gradient descent
+      in the relevant
+      direction and dampens oscillations. Defaults to 0, i.e., vanilla gradient
+      descent.
+    nesterov: boolean. Whether to apply Nesterov momentum.
+      Defaults to `False`.
+    name: Optional name prefix for the operations created when applying
+      gradients.  Defaults to `"SGD"`.
+    **kwargs: Keyword arguments. Allowed to be one of
+      `"clipnorm"` or `"clipvalue"`.
+      `"clipnorm"` (float) clips gradients by norm; `"clipvalue"` (float) clips
+      gradients by value.
+
+  Usage:
+
+  >>> opt = tf.keras.optimizers.SGD(learning_rate=0.1)
+  >>> var = tf.Variable(1.0)
+  >>> loss = lambda: (var ** 2)/2.0         # d(loss)/d(var1) = var1
+  >>> step_count = opt.minimize(loss, [var]).numpy()
+  >>> # Step is `- learning_rate * grad`
+  >>> var.numpy()
+  0.9
+
+  >>> opt = tf.keras.optimizers.SGD(learning_rate=0.1, momentum=0.9)
+  >>> var = tf.Variable(1.0)
+  >>> val0 = var.value()
+  >>> loss = lambda: (var ** 2)/2.0         # d(loss)/d(var1) = var1
+  >>> # First step is `- learning_rate * grad`
+  >>> step_count = opt.minimize(loss, [var]).numpy()
+  >>> val1 = var.value()
+  >>> (val0 - val1).numpy()
+  0.1
+  >>> # On later steps, step-size increases because of momentum
+  >>> step_count = opt.minimize(loss, [var]).numpy()
+  >>> val2 = var.value()
+  >>> (val1 - val2).numpy()
+  0.18
+
+  Reference:
+      - For `nesterov=True`, See [Sutskever et al., 2013](
+        http://jmlr.org/proceedings/papers/v28/sutskever13.pdf).
+  """
+
+  _HAS_AGGREGATE_GRAD = True
+
+  def __init__(self,
+               learning_rate=0.01,
+               momentum=0.0,
+               nesterov=False,
+               name="SGD",
+               **kwargs):
+    super(SGD, self).__init__(name, **kwargs)
+    self._set_hyper("learning_rate", kwargs.get("lr", learning_rate))
+    self._set_hyper("decay", self._initial_decay)
+
+    self._momentum = False
+    if isinstance(
+        momentum, tensor.Tensor) or callable(momentum) or momentum > 0:
+      self._momentum = True
+    if isinstance(momentum, (int, float)) and (momentum < 0 or momentum > 1):
+      raise ValueError("`momentum` must be between [0, 1].")
+    self._set_hyper("momentum", momentum)
+
+    self.nesterov = nesterov
+
+  def _create_slots(self, var_list):
+    if self._momentum:
+      for var in var_list:
+        self.add_slot(var, "momentum")
+
+  def _prepare_local(self, var_device, var_dtype, apply_state):
+    super(SGD, self)._prepare_local(var_device, var_dtype, apply_state)
+    apply_state[(var_device, var_dtype)]["momentum"] = array_ops.identity(
+        self._get_hyper("momentum", var_dtype))
+
+  def _resource_apply_dense(self, grad, var, apply_state=None):
+    var_device, var_dtype = var.device, var.dtype.base_dtype
+    coefficients = ((apply_state or {}).get((var_device, var_dtype))
+                    or self._fallback_apply_state(var_device, var_dtype))
+
+    if self._momentum:
+      momentum_var = self.get_slot(var, "momentum")
+      return gen_training_ops.ResourceApplyKerasMomentum(
+          var=var.handle,
+          accum=momentum_var.handle,
+          lr=coefficients["lr_t"],
+          grad=grad,
+          momentum=coefficients["momentum"],
+          use_locking=self._use_locking,
+          use_nesterov=self.nesterov)
+    else:
+      return gen_training_ops.ResourceApplyGradientDescent(
+          var=var.handle,
+          alpha=coefficients["lr_t"],
+          delta=grad,
+          use_locking=self._use_locking)
+
+  def _resource_apply_sparse_duplicate_indices(self, grad, var, indices,
+                                               **kwargs):
+    if self._momentum:
+      return super(SGD, self)._resource_apply_sparse_duplicate_indices(
+          grad, var, indices, **kwargs)
+    else:
+      var_device, var_dtype = var.device, var.dtype.base_dtype
+      coefficients = (kwargs.get("apply_state", {}).get((var_device, var_dtype))
+                      or self._fallback_apply_state(var_device, var_dtype))
+
+      return gen_resource_variable_ops.ResourceScatterAdd(
+          resource=var.handle,
+          indices=indices,
+          updates=-grad * coefficients["lr_t"])
+
+  def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
+    # This method is only needed for momentum optimization.
+    var_device, var_dtype = var.device, var.dtype.base_dtype
+    coefficients = ((apply_state or {}).get((var_device, var_dtype))
+                    or self._fallback_apply_state(var_device, var_dtype))
+
+    momentum_var = self.get_slot(var, "momentum")
+    return gen_training_ops.ResourceSparseApplyKerasMomentum(
+        var=var.handle,
+        accum=momentum_var.handle,
+        lr=coefficients["lr_t"],
+        grad=grad,
+        indices=indices,
+        momentum=coefficients["momentum"],
+        use_locking=self._use_locking,
+        use_nesterov=self.nesterov)
+
+  def get_config(self):
+    config = super(SGD, self).get_config()
+    config.update({
+        "learning_rate": self._serialize_hyperparameter("learning_rate"),
+        "decay": self._initial_decay,
+        "momentum": self._serialize_hyperparameter("momentum"),
+        "nesterov": self.nesterov,
+    })
+    return config

tensorflow/python/keras/optimizer_v2/legacy_learning_rate_decay.py

@@ -0,0 +1,774 @@
+# Copyright 2015 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.
+# ==============================================================================
+"""Various learning rate decay functions."""
+
+import functools
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_conversion
+from tensorflow.python.keras.optimizer_v2 import learning_rate_schedule
+from tensorflow.python.ops import math_ops
+from tensorflow.python.util import nest
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(v1=["train.exponential_decay"])
+def exponential_decay(learning_rate,
+                      global_step,
+                      decay_steps,
+                      decay_rate,
+                      staircase=False,
+                      name=None):
+  """Applies exponential decay to the learning rate.
+
+  When training a model, it is often recommended to lower the learning rate as
+  the training progresses.  This function applies an exponential decay function
+  to a provided initial learning rate.  It requires a `global_step` value to
+  compute the decayed learning rate.  You can just pass a TensorFlow variable
+  that you increment at each training step.
+
+  The function returns the decayed learning rate.  It is computed as:
+
+  ```python
+  decayed_learning_rate = learning_rate *
+                          decay_rate ^ (global_step / decay_steps)
+  ```
+
+  If the argument `staircase` is `True`, then `global_step / decay_steps` is an
+  integer division and the decayed learning rate follows a staircase function.
+
+  Example: decay every 100000 steps with a base of 0.96:
+
+  ```python
+  ...
+  global_step = tf.Variable(0, trainable=False)
+  starter_learning_rate = 0.1
+  learning_rate = tf.compat.v1.train.exponential_decay(starter_learning_rate,
+  global_step,
+                                             100000, 0.96, staircase=True)
+  # Passing global_step to minimize() will increment it at each step.
+  learning_step = (
+      tf.compat.v1.train.GradientDescentOptimizer(learning_rate)
+      .minimize(...my loss..., global_step=global_step)
+  )
+  ```
+
+  Args:
+    learning_rate: A scalar `float32` or `float64` `Tensor` or a Python number.
+      The initial learning rate.
+    global_step: A scalar `int32` or `int64` `Tensor` or a Python number. Global
+      step to use for the decay computation.  Must not be negative.
+    decay_steps: A scalar `int32` or `int64` `Tensor` or a Python number. Must
+      be positive.  See the decay computation above.
+    decay_rate: A scalar `float32` or `float64` `Tensor` or a Python number.
+      The decay rate.
+    staircase: Boolean.  If `True` decay the learning rate at discrete intervals
+    name: String.  Optional name of the operation.  Defaults to
+      'ExponentialDecay'.
+
+  Returns:
+    A scalar `Tensor` of the same type as `learning_rate`.  The decayed
+    learning rate.
+
+  Raises:
+    ValueError: if `global_step` is not supplied.
+
+  @compatibility(eager)
+  When eager execution is enabled, this function returns a function which in
+  turn returns the decayed learning rate Tensor. This can be useful for changing
+  the learning rate value across different invocations of optimizer functions.
+  @end_compatibility
+  """
+  decayed_lr = learning_rate_schedule.ExponentialDecay(
+      learning_rate, decay_steps, decay_rate, staircase=staircase, name=name)
+  if not context.executing_eagerly():
+    decayed_lr = decayed_lr(global_step)
+  else:
+    decayed_lr = functools.partial(decayed_lr, global_step)
+  return decayed_lr
+
+
+@tf_export(v1=["train.piecewise_constant_decay", "train.piecewise_constant"])
+def piecewise_constant(x, boundaries, values, name=None):
+  """Piecewise constant from boundaries and interval values.
+
+  Example: use a learning rate that's 1.0 for the first 100001 steps, 0.5
+    for the next 10000 steps, and 0.1 for any additional steps.
+
+  ```python
+  global_step = tf.Variable(0, trainable=False)
+  boundaries = [100000, 110000]
+  values = [1.0, 0.5, 0.1]
+  learning_rate = tf.compat.v1.train.piecewise_constant(global_step, boundaries,
+  values)
+
+  # Later, whenever we perform an optimization step, we increment global_step.
+  ```
+
+  Args:
+    x: A 0-D scalar `Tensor`. Must be one of the following types: `float32`,
+      `float64`, `uint8`, `int8`, `int16`, `int32`, `int64`.
+    boundaries: A list of `Tensor`s or `int`s or `float`s with strictly
+      increasing entries, and with all elements having the same type as `x`.
+    values: A list of `Tensor`s or `float`s or `int`s that specifies the values
+      for the intervals defined by `boundaries`. It should have one more element
+      than `boundaries`, and all elements should have the same type.
+    name: A string. Optional name of the operation. Defaults to
+      'PiecewiseConstant'.
+
+  Returns:
+    A 0-D Tensor. Its value is `values[0]` when `x <= boundaries[0]`,
+    `values[1]` when `x > boundaries[0]` and `x <= boundaries[1]`, ...,
+    and values[-1] when `x > boundaries[-1]`.
+
+  Raises:
+    ValueError: if types of `x` and `boundaries` do not match, or types of all
+        `values` do not match or
+        the number of elements in the lists does not match.
+
+  @compatibility(eager)
+  When eager execution is enabled, this function returns a function which in
+  turn returns the decayed learning rate Tensor. This can be useful for changing
+  the learning rate value across different invocations of optimizer functions.
+  @end_compatibility
+  """
+  boundaries = nest.map_structure(
+      tensor_conversion.convert_to_tensor_v2_with_dispatch,
+      nest.flatten(boundaries),
+  )
+  values = nest.map_structure(
+      tensor_conversion.convert_to_tensor_v2_with_dispatch, nest.flatten(values)
+  )
+  x_recomp = tensor_conversion.convert_to_tensor_v2_with_dispatch(x)
+  # Avoid explicit conversion to x's dtype. This could result in faulty
+  # comparisons, for example if floats are converted to integers.
+  for i, b in enumerate(boundaries):
+    if b.dtype.base_dtype != x_recomp.dtype.base_dtype:
+      # We can promote int32 boundaries to int64 without loss of precision.
+      # This covers the most common case where the user passes in boundaries
+      # as an array of Python integers.
+      if (b.dtype.base_dtype == dtypes.int32 and
+          x_recomp.dtype.base_dtype == dtypes.int64):
+        b = math_ops.cast(b, x_recomp.dtype.base_dtype)
+        boundaries[i] = b
+      else:
+        raise ValueError(
+            "Boundaries (%s) must have the same dtype as x (%s)." %
+            (b.dtype.base_dtype, x_recomp.dtype.base_dtype))
+  for v in values[1:]:
+    if v.dtype.base_dtype != values[0].dtype.base_dtype:
+      raise ValueError(
+          "Values must have elements all with the same dtype (%s vs %s)." %
+          (values[0].dtype.base_dtype, v.dtype.base_dtype))
+  decayed_lr = learning_rate_schedule.PiecewiseConstantDecay(
+      boundaries, values, name=name)
+  if not context.executing_eagerly():
+    decayed_lr = decayed_lr(x)
+  else:
+    decayed_lr = functools.partial(decayed_lr, x)
+  return decayed_lr
+
+
+@tf_export(v1=["train.polynomial_decay"])
+def polynomial_decay(learning_rate,
+                     global_step,
+                     decay_steps,
+                     end_learning_rate=0.0001,
+                     power=1.0,
+                     cycle=False,
+                     name=None):
+  """Applies a polynomial decay to the learning rate.
+
+  It is commonly observed that a monotonically decreasing learning rate, whose
+  degree of change is carefully chosen, results in a better performing model.
+  This function applies a polynomial decay function to a provided initial
+  `learning_rate` to reach an `end_learning_rate` in the given `decay_steps`.
+
+  It requires a `global_step` value to compute the decayed learning rate.  You
+  can just pass a TensorFlow variable that you increment at each training step.
+
+  The function returns the decayed learning rate.  It is computed as:
+
+  ```python
+  global_step = min(global_step, decay_steps)
+  decayed_learning_rate = (learning_rate - end_learning_rate) *
+                          (1 - global_step / decay_steps) ^ (power) +
+                          end_learning_rate
+
+  ```
+
+  If `cycle` is True then a multiple of `decay_steps` is used, the first one
+  that is bigger than `global_steps`.
+
+  ```python
+  decay_steps = decay_steps * ceil(global_step / decay_steps)
+  decayed_learning_rate = (learning_rate - end_learning_rate) *
+                          (1 - global_step / decay_steps) ^ (power) +
+                          end_learning_rate
+
+  ```
+
+  Example: decay from 0.1 to 0.01 in 10000 steps using sqrt (i.e. power=0.5):
+
+  ```python
+  ...
+  global_step = tf.Variable(0, trainable=False)
+  starter_learning_rate = 0.1
+  end_learning_rate = 0.01
+  decay_steps = 10000
+  learning_rate = tf.compat.v1.train.polynomial_decay(starter_learning_rate,
+  global_step,
+                                            decay_steps, end_learning_rate,
+                                            power=0.5)
+  # Passing global_step to minimize() will increment it at each step.
+  learning_step = (
+      tf.compat.v1.train.GradientDescentOptimizer(learning_rate)
+      .minimize(...my loss..., global_step=global_step)
+  )
+  ```
+
+  Args:
+    learning_rate: A scalar `float32` or `float64` `Tensor` or a Python number.
+      The initial learning rate.
+    global_step: A scalar `int32` or `int64` `Tensor` or a Python number. Global
+      step to use for the decay computation.  Must not be negative.
+    decay_steps: A scalar `int32` or `int64` `Tensor` or a Python number. Must
+      be positive.  See the decay computation above.
+    end_learning_rate: A scalar `float32` or `float64` `Tensor` or a Python
+      number.  The minimal end learning rate.
+    power: A scalar `float32` or `float64` `Tensor` or a Python number.  The
+      power of the polynomial. Defaults to linear, 1.0.
+    cycle: A boolean, whether or not it should cycle beyond decay_steps.
+    name: String.  Optional name of the operation. Defaults to
+      'PolynomialDecay'.
+
+  Returns:
+    A scalar `Tensor` of the same type as `learning_rate`.  The decayed
+    learning rate.
+
+  Raises:
+    ValueError: if `global_step` is not supplied.
+
+  @compatibility(eager)
+  When eager execution is enabled, this function returns a function which in
+  turn returns the decayed learning rate Tensor. This can be useful for changing
+  the learning rate value across different invocations of optimizer functions.
+  @end_compatibility
+  """
+  decayed_lr = learning_rate_schedule.PolynomialDecay(
+      learning_rate,
+      decay_steps,
+      end_learning_rate=end_learning_rate,
+      power=power,
+      cycle=cycle,
+      name=name)
+
+  if not context.executing_eagerly():
+    decayed_lr = decayed_lr(global_step)
+  else:
+    decayed_lr = functools.partial(decayed_lr, global_step)
+  return decayed_lr
+
+
+@tf_export(v1=["train.natural_exp_decay"])
+def natural_exp_decay(learning_rate,
+                      global_step,
+                      decay_steps,
+                      decay_rate,
+                      staircase=False,
+                      name=None):
+  """Applies natural exponential decay to the initial learning rate.
+
+  When training a model, it is often recommended to lower the learning rate as
+  the training progresses.  This function applies an exponential decay function
+  to a provided initial learning rate.  It requires an `global_step` value to
+  compute the decayed learning rate.  You can just pass a TensorFlow variable
+  that you increment at each training step.
+
+  The function returns the decayed learning rate.  It is computed as:
+
+  ```python
+  decayed_learning_rate = learning_rate * exp(-decay_rate * global_step /
+  decay_step)
+  ```
+
+  or, if `staircase` is `True`, as:
+
+  ```python
+  decayed_learning_rate = learning_rate * exp(-decay_rate * floor(global_step /
+  decay_step))
+  ```
+
+  Example: decay exponentially with a base of 0.96:
+
+  ```python
+  ...
+  global_step = tf.Variable(0, trainable=False)
+  learning_rate = 0.1
+  decay_steps = 5
+  k = 0.5
+  learning_rate = tf.compat.v1.train.natural_exp_decay(learning_rate,
+  global_step,
+                                             decay_steps, k)
+
+  # Passing global_step to minimize() will increment it at each step.
+  learning_step = (
+      tf.compat.v1.train.GradientDescentOptimizer(learning_rate)
+      .minimize(...my loss..., global_step=global_step)
+  )
+  ```
+
+  Args:
+    learning_rate: A scalar `float32` or `float64` `Tensor` or a Python number.
+      The initial learning rate.
+    global_step: A Python number. Global step to use for the decay computation.
+      Must not be negative.
+    decay_steps: How often to apply decay.
+    decay_rate: A Python number.  The decay rate.
+    staircase: Whether to apply decay in a discrete staircase, as opposed to
+      continuous, fashion.
+    name: String.  Optional name of the operation.  Defaults to
+      'ExponentialTimeDecay'.
+
+  Returns:
+    A scalar `Tensor` of the same type as `learning_rate`.  The decayed
+    learning rate.
+
+  Raises:
+    ValueError: if `global_step` is not supplied.
+
+  @compatibility(eager)
+  When eager execution is enabled, this function returns a function which in
+  turn returns the decayed learning rate Tensor. This can be useful for changing
+  the learning rate value across different invocations of optimizer functions.
+  @end_compatibility
+  """
+  natural_exp_rate = math_ops.exp(math_ops.negative(decay_rate))
+  decayed_lr = learning_rate_schedule.ExponentialDecay(
+      learning_rate,
+      decay_steps,
+      natural_exp_rate,
+      staircase=staircase,
+      name=name)
+
+  if not context.executing_eagerly():
+    decayed_lr = decayed_lr(global_step)
+  else:
+    decayed_lr = functools.partial(decayed_lr, global_step)
+  return decayed_lr
+
+
+@tf_export(v1=["train.inverse_time_decay"])
+def inverse_time_decay(learning_rate,
+                       global_step,
+                       decay_steps,
+                       decay_rate,
+                       staircase=False,
+                       name=None):
+  """Applies inverse time decay to the initial learning rate.
+
+  When training a model, it is often recommended to lower the learning rate as
+  the training progresses.  This function applies an inverse decay function
+  to a provided initial learning rate.  It requires an `global_step` value to
+  compute the decayed learning rate.  You can just pass a TensorFlow variable
+  that you increment at each training step.
+
+  The function returns the decayed learning rate.  It is computed as:
+
+  ```python
+  decayed_learning_rate = learning_rate / (1 + decay_rate * global_step /
+  decay_step)
+  ```
+
+  or, if `staircase` is `True`, as:
+
+  ```python
+  decayed_learning_rate = learning_rate / (1 + decay_rate * floor(global_step /
+  decay_step))
+  ```
+
+  Example: decay 1/t with a rate of 0.5:
+
+  ```python
+  ...
+  global_step = tf.Variable(0, trainable=False)
+  learning_rate = 0.1
+  decay_steps = 1.0
+  decay_rate = 0.5
+  learning_rate = tf.compat.v1.train.inverse_time_decay(learning_rate,
+  global_step,
+  decay_steps, decay_rate)
+
+  # Passing global_step to minimize() will increment it at each step.
+  learning_step = (
+      tf.compat.v1.train.GradientDescentOptimizer(learning_rate)
+      .minimize(...my loss..., global_step=global_step)
+  )
+  ```
+
+  Args:
+    learning_rate: A scalar `float32` or `float64` `Tensor` or a Python number.
+      The initial learning rate.
+    global_step: A Python number. Global step to use for the decay computation.
+      Must not be negative.
+    decay_steps: How often to apply decay.
+    decay_rate: A Python number.  The decay rate.
+    staircase: Whether to apply decay in a discrete staircase, as opposed to
+      continuous, fashion.
+    name: String.  Optional name of the operation.  Defaults to
+      'InverseTimeDecay'.
+
+  Returns:
+    A scalar `Tensor` of the same type as `learning_rate`.  The decayed
+    learning rate.
+
+  Raises:
+    ValueError: if `global_step` is not supplied.
+
+  @compatibility(eager)
+  When eager execution is enabled, this function returns a function which in
+  turn returns the decayed learning rate Tensor. This can be useful for changing
+  the learning rate value across different invocations of optimizer functions.
+  @end_compatibility
+  """
+  decayed_lr = learning_rate_schedule.InverseTimeDecay(
+      learning_rate, decay_steps, decay_rate, staircase=staircase, name=name)
+
+  if not context.executing_eagerly():
+    decayed_lr = decayed_lr(global_step)
+  else:
+    decayed_lr = functools.partial(decayed_lr, global_step)
+  return decayed_lr
+
+
+@tf_export(v1=["train.cosine_decay"])
+def cosine_decay(learning_rate, global_step, decay_steps, alpha=0.0, name=None):
+  """Applies cosine decay to the learning rate.
+
+  When training a model, it is often recommended to lower the learning rate as
+  the training progresses.  This function applies a cosine decay function
+  to a provided initial learning rate.  It requires a `global_step` value to
+  compute the decayed learning rate.  You can just pass a TensorFlow variable
+  that you increment at each training step.
+
+  The function returns the decayed learning rate.  It is computed as:
+  ```python
+  global_step = min(global_step, decay_steps)
+  cosine_decay = 0.5 * (1 + cos(pi * global_step / decay_steps))
+  decayed = (1 - alpha) * cosine_decay + alpha
+  decayed_learning_rate = learning_rate * decayed
+  ```
+
+  Example usage:
+  ```python
+  decay_steps = 1000
+  lr_decayed = cosine_decay(learning_rate, global_step, decay_steps)
+  ```
+
+  Args:
+    learning_rate: A scalar `float32` or `float64` Tensor or a Python number.
+      The initial learning rate.
+    global_step: A scalar `int32` or `int64` `Tensor` or a Python number. Global
+      step to use for the decay computation.
+    decay_steps: A scalar `int32` or `int64` `Tensor` or a Python number. Number
+      of steps to decay over.
+    alpha: A scalar `float32` or `float64` Tensor or a Python number. Minimum
+      learning rate value as a fraction of learning_rate.
+    name: String. Optional name of the operation.  Defaults to 'CosineDecay'.
+
+  Returns:
+    A scalar `Tensor` of the same type as `learning_rate`.  The decayed
+    learning rate.
+  Raises:
+    ValueError: if `global_step` is not supplied.
+
+  References:
+    Stochastic Gradient Descent with Warm Restarts:
+      [Loshchilov et al., 2017]
+      (https://openreview.net/forum?id=Skq89Scxx&noteId=Skq89Scxx)
+      ([pdf](https://openreview.net/pdf?id=Skq89Scxx))
+
+  @compatibility(eager)
+  When eager execution is enabled, this function returns a function which in
+  turn returns the decayed learning rate Tensor. This can be useful for changing
+  the learning rate value across different invocations of optimizer functions.
+  @end_compatibility
+  """
+  decayed_lr = learning_rate_schedule.CosineDecay(
+      learning_rate, decay_steps, alpha=alpha, name=name)
+
+  if not context.executing_eagerly():
+    decayed_lr = decayed_lr(global_step)
+  else:
+    decayed_lr = functools.partial(decayed_lr, global_step)
+  return decayed_lr
+
+
+@tf_export(v1=["train.cosine_decay_restarts"])
+def cosine_decay_restarts(learning_rate,
+                          global_step,
+                          first_decay_steps,
+                          t_mul=2.0,
+                          m_mul=1.0,
+                          alpha=0.0,
+                          name=None):
+  """Applies cosine decay with restarts to the learning rate.
+
+  When training a model, it is often recommended to lower the learning rate as
+  the training progresses.  This function applies a cosine decay function with
+  restarts to a provided initial learning rate.  It requires a `global_step`
+  value to compute the decayed learning rate.  You can just pass a TensorFlow
+  variable that you increment at each training step.
+
+  The function returns the decayed learning rate while taking into account
+  possible warm restarts. The learning rate multiplier first decays
+  from 1 to `alpha` for `first_decay_steps` steps. Then, a warm
+  restart is performed. Each new warm restart runs for `t_mul` times more steps
+  and with `m_mul` times smaller initial learning rate.
+
+  Example usage:
+  ```python
+  first_decay_steps = 1000
+  lr_decayed = cosine_decay_restarts(learning_rate, global_step,
+                                     first_decay_steps)
+  ```
+
+  Args:
+    learning_rate: A scalar `float32` or `float64` Tensor or a Python number.
+      The initial learning rate.
+    global_step: A scalar `int32` or `int64` `Tensor` or a Python number. Global
+      step to use for the decay computation.
+    first_decay_steps: A scalar `int32` or `int64` `Tensor` or a Python number.
+      Number of steps to decay over.
+    t_mul: A scalar `float32` or `float64` `Tensor` or a Python number. Used to
+      derive the number of iterations in the i-th period
+    m_mul: A scalar `float32` or `float64` `Tensor` or a Python number.
+      Used to derive the initial learning rate of the i-th period:
+    alpha: A scalar `float32` or `float64` Tensor or a Python number. Minimum
+      learning rate value as a fraction of the learning_rate.
+    name: String. Optional name of the operation.  Defaults to 'SGDRDecay'.
+
+  Returns:
+    A scalar `Tensor` of the same type as `learning_rate`.  The decayed
+    learning rate.
+  Raises:
+    ValueError: if `global_step` is not supplied.
+
+  References:
+    Stochastic Gradient Descent with Warm Restarts:
+      [Loshchilov et al., 2017]
+      (https://openreview.net/forum?id=Skq89Scxx&noteId=Skq89Scxx)
+      ([pdf](https://openreview.net/pdf?id=Skq89Scxx))
+
+  @compatibility(eager)
+  When eager execution is enabled, this function returns a function which in
+  turn returns the decayed learning rate Tensor. This can be useful for changing
+  the learning rate value across different invocations of optimizer functions.
+  @end_compatibility
+  """
+  decayed_lr = learning_rate_schedule.CosineDecayRestarts(
+      learning_rate,
+      first_decay_steps,
+      t_mul=t_mul,
+      m_mul=m_mul,
+      alpha=alpha,
+      name=name)
+
+  if not context.executing_eagerly():
+    decayed_lr = decayed_lr(global_step)
+  else:
+    decayed_lr = functools.partial(decayed_lr, global_step)
+  return decayed_lr
+
+
+@tf_export(v1=["train.linear_cosine_decay"])
+def linear_cosine_decay(learning_rate,
+                        global_step,
+                        decay_steps,
+                        num_periods=0.5,
+                        alpha=0.0,
+                        beta=0.001,
+                        name=None):
+  """Applies linear cosine decay to the learning rate.
+
+  Note that linear cosine decay is more aggressive than cosine decay and
+  larger initial learning rates can typically be used.
+
+  When training a model, it is often recommended to lower the learning rate as
+  the training progresses.  This function applies a linear cosine decay function
+  to a provided initial learning rate.  It requires a `global_step` value to
+  compute the decayed learning rate.  You can just pass a TensorFlow variable
+  that you increment at each training step.
+
+  The function returns the decayed learning rate.  It is computed as:
+  ```python
+  global_step = min(global_step, decay_steps)
+  linear_decay = (decay_steps - global_step) / decay_steps)
+  cosine_decay = 0.5 * (
+      1 + cos(pi * 2 * num_periods * global_step / decay_steps))
+  decayed = (alpha + linear_decay) * cosine_decay + beta
+  decayed_learning_rate = learning_rate * decayed
+  ```
+
+  Example usage:
+  ```python
+  decay_steps = 1000
+  lr_decayed = linear_cosine_decay(learning_rate, global_step, decay_steps)
+  ```
+
+  Args:
+    learning_rate: A scalar `float32` or `float64` Tensor or a Python number.
+      The initial learning rate.
+    global_step: A scalar `int32` or `int64` `Tensor` or a Python number. Global
+      step to use for the decay computation.
+    decay_steps: A scalar `int32` or `int64` `Tensor` or a Python number. Number
+      of steps to decay over.
+    num_periods: Number of periods in the cosine part of the decay. See
+      computation above.
+    alpha: See computation above.
+    beta: See computation above.
+    name: String.  Optional name of the operation.  Defaults to
+      'LinearCosineDecay'.
+
+  Returns:
+    A scalar `Tensor` of the same type as `learning_rate`.  The decayed
+    learning rate.
+  Raises:
+    ValueError: if `global_step` is not supplied.
+
+  References:
+    Neural Optimizer Search with Reinforcement Learning:
+      [Bello et al., 2017](http://proceedings.mlr.press/v70/bello17a.html)
+      ([pdf](http://proceedings.mlr.press/v70/bello17a/bello17a.pdf))
+    Stochastic Gradient Descent with Warm Restarts:
+      [Loshchilov et al., 2017]
+      (https://openreview.net/forum?id=Skq89Scxx&noteId=Skq89Scxx)
+      ([pdf](https://openreview.net/pdf?id=Skq89Scxx))
+
+  @compatibility(eager)
+  When eager execution is enabled, this function returns a function which in
+  turn returns the decayed learning rate Tensor. This can be useful for changing
+  the learning rate value across different invocations of optimizer functions.
+  @end_compatibility
+  """
+  decayed_lr = learning_rate_schedule.LinearCosineDecay(
+      learning_rate,
+      decay_steps,
+      num_periods=num_periods,
+      alpha=alpha,
+      beta=beta,
+      name=name)
+
+  if not context.executing_eagerly():
+    decayed_lr = decayed_lr(global_step)
+  else:
+    decayed_lr = functools.partial(decayed_lr, global_step)
+  return decayed_lr
+
+
+@tf_export(v1=["train.noisy_linear_cosine_decay"])
+def noisy_linear_cosine_decay(learning_rate,
+                              global_step,
+                              decay_steps,
+                              initial_variance=1.0,
+                              variance_decay=0.55,
+                              num_periods=0.5,
+                              alpha=0.0,
+                              beta=0.001,
+                              name=None):
+  """Applies noisy linear cosine decay to the learning rate.
+
+  Note that linear cosine decay is more aggressive than cosine decay and
+  larger initial learning rates can typically be used.
+
+  When training a model, it is often recommended to lower the learning rate as
+  the training progresses.  This function applies a noisy linear
+  cosine decay function to a provided initial learning rate.
+  It requires a `global_step` value to compute the decayed learning rate.
+  You can just pass a TensorFlow variable that you increment at each
+  training step.
+
+  The function returns the decayed learning rate.  It is computed as:
+  ```python
+  global_step = min(global_step, decay_steps)
+  linear_decay = (decay_steps - global_step) / decay_steps)
+  cosine_decay = 0.5 * (
+      1 + cos(pi * 2 * num_periods * global_step / decay_steps))
+  decayed = (alpha + linear_decay + eps_t) * cosine_decay + beta
+  decayed_learning_rate = learning_rate * decayed
+  ```
+  where eps_t is 0-centered gaussian noise with variance
+  initial_variance / (1 + global_step) ** variance_decay
+
+  Example usage:
+  ```python
+  decay_steps = 1000
+  lr_decayed = noisy_linear_cosine_decay(
+    learning_rate, global_step, decay_steps)
+  ```
+
+  Args:
+    learning_rate: A scalar `float32` or `float64` Tensor or a Python number.
+      The initial learning rate.
+    global_step: A scalar `int32` or `int64` `Tensor` or a Python number. Global
+      step to use for the decay computation.
+    decay_steps: A scalar `int32` or `int64` `Tensor` or a Python number. Number
+      of steps to decay over.
+    initial_variance: initial variance for the noise. See computation above.
+    variance_decay: decay for the noise's variance. See computation above.
+    num_periods: Number of periods in the cosine part of the decay. See
+      computation above.
+    alpha: See computation above.
+    beta: See computation above.
+    name: String.  Optional name of the operation.  Defaults to
+      'NoisyLinearCosineDecay'.
+
+  Returns:
+    A scalar `Tensor` of the same type as `learning_rate`.  The decayed
+    learning rate.
+  Raises:
+    ValueError: if `global_step` is not supplied.
+
+  References:
+    Neural Optimizer Search with Reinforcement Learning:
+      [Bello et al., 2017](http://proceedings.mlr.press/v70/bello17a.html)
+      ([pdf](http://proceedings.mlr.press/v70/bello17a/bello17a.pdf))
+    Stochastic Gradient Descent with Warm Restarts:
+      [Loshchilov et al., 2017]
+      (https://openreview.net/forum?id=Skq89Scxx&noteId=Skq89Scxx)
+      ([pdf](https://openreview.net/pdf?id=Skq89Scxx))
+
+  @compatibility(eager)
+  When eager execution is enabled, this function returns a function which in
+  turn returns the decayed learning rate Tensor. This can be useful for changing
+  the learning rate value across different invocations of optimizer functions.
+  @end_compatibility
+  """
+  decayed_lr = learning_rate_schedule.NoisyLinearCosineDecay(
+      learning_rate,
+      decay_steps,
+      initial_variance=initial_variance,
+      variance_decay=variance_decay,
+      num_periods=num_periods,
+      alpha=alpha,
+      beta=beta,
+      name=name)
+
+  if not context.executing_eagerly():
+    decayed_lr = decayed_lr(global_step)
+  else:
+    decayed_lr = functools.partial(decayed_lr, global_step)
+  return decayed_lr

tensorflow/python/keras/optimizer_v2/nadam.py

@@ -0,0 +1,220 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Nadam optimizer implementation."""
+# pylint: disable=g-classes-have-attributes
+
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_conversion
+from tensorflow.python.keras import backend_config
+from tensorflow.python.keras.optimizer_v2 import learning_rate_schedule
+from tensorflow.python.keras.optimizer_v2 import optimizer_v2
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import state_ops
+from tensorflow.python.ops import variables as tf_variables
+
+
+class Nadam(optimizer_v2.OptimizerV2):
+  r"""Optimizer that implements the NAdam algorithm.
+  Much like Adam is essentially RMSprop with momentum, Nadam is Adam with
+  Nesterov momentum.
+
+  Args:
+    learning_rate: A Tensor or a floating point value.  The learning rate.
+    beta_1: A float value or a constant float tensor. The exponential decay
+      rate for the 1st moment estimates.
+    beta_2: A float value or a constant float tensor. The exponential decay
+      rate for the exponentially weighted infinity norm.
+    epsilon: A small constant for numerical stability.
+    name: Optional name for the operations created when applying gradients.
+      Defaults to `"Nadam"`.
+    **kwargs: Keyword arguments. Allowed to be one of
+      `"clipnorm"` or `"clipvalue"`.
+      `"clipnorm"` (float) clips gradients by norm; `"clipvalue"` (float) clips
+      gradients by value.
+
+  Usage Example:
+    >>> opt = tf.keras.optimizers.Nadam(learning_rate=0.2)
+    >>> var1 = tf.Variable(10.0)
+    >>> loss = lambda: (var1 ** 2) / 2.0
+    >>> step_count = opt.minimize(loss, [var1]).numpy()
+    >>> "{:.1f}".format(var1.numpy())
+    9.8
+
+  Reference:
+    - [Dozat, 2015](http://cs229.stanford.edu/proj2015/054_report.pdf).
+  """
+
+  _HAS_AGGREGATE_GRAD = True
+
+  def __init__(self,
+               learning_rate=0.001,
+               beta_1=0.9,
+               beta_2=0.999,
+               epsilon=1e-7,
+               name='Nadam',
+               **kwargs):
+    # Backwards compatibility with keras NAdam optimizer.
+    kwargs['decay'] = kwargs.pop('schedule_decay', 0.004)
+    learning_rate = kwargs.get('lr', learning_rate)
+    if isinstance(learning_rate, learning_rate_schedule.LearningRateSchedule):
+      raise ValueError('The Nadam optimizer does not support '
+                       'tf.keras.optimizers.LearningRateSchedules as the '
+                       'learning rate.')
+
+    super(Nadam, self).__init__(name, **kwargs)
+    self._set_hyper('learning_rate', kwargs.get('lr', learning_rate))
+    self._set_hyper('decay', self._initial_decay)
+    self._set_hyper('beta_1', beta_1)
+    self._set_hyper('beta_2', beta_2)
+    self.epsilon = epsilon or backend_config.epsilon()
+    self._m_cache = None
+
+  def _create_slots(self, var_list):
+    var_dtype = var_list[0].dtype.base_dtype
+    if self._m_cache is None:
+      self._m_cache = self.add_weight(
+          'momentum_cache',
+          shape=[],
+          dtype=var_dtype,
+          initializer='ones',
+          trainable=False,
+          aggregation=tf_variables.VariableAggregation.ONLY_FIRST_REPLICA)
+      self._weights.append(self._m_cache)
+    # Separate for-loops to respect the ordering of slot variables from v1.
+    for var in var_list:
+      # Create slots for the first moments.
+      self.add_slot(var, 'm')
+    for var in var_list:
+      # Create slots for the second moments.
+      self.add_slot(var, 'v')
+
+  def _prepare_local(self, var_device, var_dtype, apply_state):
+    lr_t = array_ops.identity(self._get_hyper('learning_rate', var_dtype))
+    beta_1_t = array_ops.identity(self._get_hyper('beta_1', var_dtype))
+    beta_2_t = array_ops.identity(self._get_hyper('beta_2', var_dtype))
+    local_step = math_ops.cast(self.iterations + 1, var_dtype)
+    next_step = math_ops.cast(self.iterations + 2, var_dtype)
+
+    decay_base = math_ops.cast(0.96, var_dtype)
+
+    m_t = beta_1_t * (1. - 0.5 * (
+        math_ops.pow(decay_base, self._initial_decay * local_step)))
+    m_t_1 = beta_1_t * (1. - 0.5 * (
+        math_ops.pow(decay_base, self._initial_decay * next_step)))
+
+    m_schedule_new = math_ops.cast(self._m_cache_read, var_dtype) * m_t
+    if var_dtype is self._m_cache.dtype:
+      m_schedule_new = array_ops.identity(state_ops.assign(
+          self._m_cache, m_schedule_new, use_locking=self._use_locking))
+    m_schedule_next = m_schedule_new * m_t_1
+
+    apply_state[(var_device, var_dtype)] = dict(
+        lr_t=lr_t,
+        neg_lr_t=-lr_t,  # pylint: disable=invalid-unary-operand-type
+        epsilon=tensor_conversion.convert_to_tensor_v2_with_dispatch(
+            self.epsilon, var_dtype
+        ),
+        beta_1_t=beta_1_t,
+        beta_2_t=beta_2_t,
+        m_t=m_t,
+        m_t_1=m_t_1,
+        one_minus_beta_1_t=1 - beta_1_t,
+        one_minus_beta_2_t=1 - beta_2_t,
+        one_minus_m_t=1.0 - m_t,
+        one_minus_m_schedule_new=1.0 - m_schedule_new,
+        one_minus_m_schedule_next=1.0 - m_schedule_next,
+        v_t_prime_denominator=1.0 - math_ops.pow(beta_2_t, local_step),
+    )
+
+  def _prepare(self, var_list):
+    # Get the value of the momentum cache before starting to apply gradients.
+    self._m_cache_read = array_ops.identity(self._m_cache)
+    return super(Nadam, self)._prepare(var_list)
+
+  def _resource_apply_dense(self, grad, var, apply_state=None):
+    var_device, var_dtype = var.device, var.dtype.base_dtype
+    coefficients = ((apply_state or {}).get((var_device, var_dtype))
+                    or self._fallback_apply_state(var_device, var_dtype))
+
+    m = self.get_slot(var, 'm')
+    v = self.get_slot(var, 'v')
+
+    g_prime = grad / coefficients['one_minus_m_schedule_new']
+    m_t = (coefficients['beta_1_t'] * m +
+           coefficients['one_minus_beta_1_t'] * grad)
+    m_t = state_ops.assign(m, m_t, use_locking=self._use_locking)
+    m_t_prime = m_t / coefficients['one_minus_m_schedule_next']
+    v_t = (coefficients['beta_2_t'] * v +
+           coefficients['one_minus_beta_2_t'] * math_ops.square(grad))
+    v_t = state_ops.assign(v, v_t, use_locking=self._use_locking)
+    v_t_prime = v_t / coefficients['v_t_prime_denominator']
+    m_t_bar = (coefficients['one_minus_m_t'] * g_prime +
+               coefficients['m_t_1'] * m_t_prime)
+    var_t = var - coefficients['lr_t'] * m_t_bar / (
+        math_ops.sqrt(v_t_prime) + coefficients['epsilon'])
+    return state_ops.assign(var, var_t, use_locking=self._use_locking).op
+
+  def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
+    var_device, var_dtype = var.device, var.dtype.base_dtype
+    coefficients = ((apply_state or {}).get((var_device, var_dtype))
+                    or self._fallback_apply_state(var_device, var_dtype))
+
+    m = self.get_slot(var, 'm')
+    v = self.get_slot(var, 'v')
+
+    g_prime = grad / coefficients['one_minus_m_schedule_new']
+
+    # m_t = beta1 * m + (1 - beta1) * g_t
+    m_scaled_g_values = grad * coefficients['one_minus_beta_1_t']
+    m_t = state_ops.assign(m, m * coefficients['beta_1_t'],
+                           use_locking=self._use_locking)
+
+    with ops.control_dependencies([m_t]):
+      m_t = self._resource_scatter_add(m, indices, m_scaled_g_values)
+      m_t_slice = array_ops.gather(m_t, indices)
+
+    m_t_prime = m_t_slice / coefficients['one_minus_m_schedule_next']
+    m_t_bar = (coefficients['one_minus_m_t'] * g_prime +
+               coefficients['m_t_1'] * m_t_prime)
+
+    # v_t = beta2 * v + (1 - beta2) * (g_t * g_t)
+    v_scaled_g_values = (grad * grad) * coefficients['one_minus_beta_2_t']
+    v_t = state_ops.assign(v, v * coefficients['beta_2_t'],
+                           use_locking=self._use_locking)
+
+    with ops.control_dependencies([v_t]):
+      v_t = self._resource_scatter_add(v, indices, v_scaled_g_values)
+      v_t_slice = array_ops.gather(v_t, indices)
+
+    v_t_prime = v_t_slice / coefficients['v_t_prime_denominator']
+    v_prime_sqrt_plus_eps = math_ops.sqrt(v_t_prime) + coefficients['epsilon']
+
+    var_update = self._resource_scatter_add(
+        var, indices,
+        coefficients['neg_lr_t'] * m_t_bar / v_prime_sqrt_plus_eps)
+    return control_flow_ops.group(*[var_update, m_t_bar, v_t])
+
+  def get_config(self):
+    config = super(Nadam, self).get_config()
+    config.update({
+        'learning_rate': self._serialize_hyperparameter('learning_rate'),
+        'decay': self._initial_decay,
+        'beta_1': self._serialize_hyperparameter('beta_1'),
+        'beta_2': self._serialize_hyperparameter('beta_2'),
+        'epsilon': self.epsilon,
+    })
+    return config

tensorflow/python/keras/optimizer_v2/rmsprop.py

@@ -0,0 +1,303 @@
+# Copyright 2020 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."""
+# pylint: disable=g-classes-have-attributes
+
+import numpy as np
+
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework import tensor_conversion
+from tensorflow.python.keras import backend_config
+from tensorflow.python.keras.optimizer_v2 import optimizer_v2
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import gen_training_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import state_ops
+
+
+class RMSprop(optimizer_v2.OptimizerV2):
+  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.
+
+  Args:
+    learning_rate: A `Tensor`, floating point value, or a schedule that is a
+      `tf.keras.optimizers.schedules.LearningRateSchedule`, or a callable
+      that takes no arguments and returns the actual value to use. The
+      learning rate. Defaults to 0.001.
+    rho: Discounting factor for the history/coming gradient. Defaults to 0.9.
+    momentum: A scalar or a scalar `Tensor`. Defaults to 0.0.
+    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`.
+    name: Optional name prefix for the operations created when applying
+      gradients. Defaults to `"RMSprop"`.
+    **kwargs: Keyword arguments. Allowed to be one of
+      `"clipnorm"` or `"clipvalue"`.
+      `"clipnorm"` (float) clips gradients by norm; `"clipvalue"` (float) clips
+      gradients by value.
+
+  Note that in the dense implementation of this algorithm, variables and their
+  corresponding accumulators (momentum, gradient moving average, square
+  gradient moving average) will be updated even if the gradient is zero
+  (i.e. accumulators will decay, momentum will be applied). The sparse
+  implementation (used when the gradient is an `IndexedSlices` object,
+  typically because of `tf.gather` or an embedding lookup in the forward pass)
+  will not update variable slices or their accumulators unless those slices
+  were used in the forward pass (nor is there an "eventual" correction to
+  account for these omitted updates). This leads to more efficient updates for
+  large embedding lookup tables (where most of the slices are not accessed in
+  a particular graph execution), but differs from the published algorithm.
+
+  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)
+  """
+
+  _HAS_AGGREGATE_GRAD = True
+
+  def __init__(self,
+               learning_rate=0.001,
+               rho=0.9,
+               momentum=0.0,
+               epsilon=1e-7,
+               centered=False,
+               name="RMSprop",
+               **kwargs):
+    """Construct a new RMSprop optimizer.
+
+    Args:
+      learning_rate: A `Tensor`, floating point value, or a schedule that is a
+        `tf.keras.optimizers.schedules.LearningRateSchedule`, or a callable
+        that takes no arguments and returns the actual value to use. The
+        learning rate. Defaults to 0.001.
+      rho: Discounting factor for the history/coming gradient. Defaults to 0.9.
+      momentum: A scalar or a scalar `Tensor`. Defaults to 0.0.
+      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`.
+      name: Optional name prefix for the operations created when applying
+        gradients. Defaults to "RMSprop".
+      **kwargs: keyword arguments. Allowed to be {`clipnorm`, `clipvalue`, `lr`,
+        `decay`}. `clipnorm` is clip gradients by norm; `clipvalue` is clip
+        gradients by value, `decay` is included for backward compatibility to
+        allow time inverse decay of learning rate. `lr` is included for backward
+        compatibility, recommended to use `learning_rate` instead.
+
+    @compatibility(eager)
+    When eager execution is enabled, `learning_rate`, `decay`, `momentum`, and
+    `epsilon` can each be a callable that takes no arguments and returns the
+    actual value to use. This can be useful for changing these values across
+    different invocations of optimizer functions.
+    @end_compatibility
+    """
+    super(RMSprop, self).__init__(name, **kwargs)
+    self._set_hyper("learning_rate", kwargs.get("lr", learning_rate))
+    self._set_hyper("decay", self._initial_decay)
+    self._set_hyper("rho", rho)
+
+    self._momentum = False
+    if isinstance(
+        momentum, tensor.Tensor) or callable(momentum) or momentum > 0:
+      self._momentum = True
+    if isinstance(momentum, (int, float)) and (momentum < 0 or momentum > 1):
+      raise ValueError("`momentum` must be between [0, 1].")
+    self._set_hyper("momentum", momentum)
+
+    self.epsilon = epsilon or backend_config.epsilon()
+    self.centered = centered
+
+  def _create_slots(self, var_list):
+    for var in var_list:
+      self.add_slot(var, "rms")
+    if self._momentum:
+      for var in var_list:
+        self.add_slot(var, "momentum")
+    if self.centered:
+      for var in var_list:
+        self.add_slot(var, "mg")
+
+  def _prepare_local(self, var_device, var_dtype, apply_state):
+    super(RMSprop, self)._prepare_local(var_device, var_dtype, apply_state)
+
+    rho = array_ops.identity(self._get_hyper("rho", var_dtype))
+    apply_state[(var_device, var_dtype)].update(
+        dict(
+            neg_lr_t=-apply_state[(var_device, var_dtype)]["lr_t"],
+            epsilon=tensor_conversion.convert_to_tensor_v2_with_dispatch(
+                self.epsilon, var_dtype
+            ),
+            rho=rho,
+            momentum=array_ops.identity(self._get_hyper("momentum", var_dtype)),
+            one_minus_rho=1.0 - rho,
+        )
+    )
+
+  def _resource_apply_dense(self, grad, var, apply_state=None):
+    var_device, var_dtype = var.device, var.dtype.base_dtype
+    coefficients = ((apply_state or {}).get((var_device, var_dtype))
+                    or self._fallback_apply_state(var_device, var_dtype))
+
+    rms = self.get_slot(var, "rms")
+    if self._momentum:
+      mom = self.get_slot(var, "momentum")
+      if self.centered:
+        mg = self.get_slot(var, "mg")
+        return gen_training_ops.ResourceApplyCenteredRMSProp(
+            var=var.handle,
+            mg=mg.handle,
+            ms=rms.handle,
+            mom=mom.handle,
+            lr=coefficients["lr_t"],
+            rho=coefficients["rho"],
+            momentum=coefficients["momentum"],
+            epsilon=coefficients["epsilon"],
+            grad=grad,
+            use_locking=self._use_locking)
+      else:
+        return gen_training_ops.ResourceApplyRMSProp(
+            var=var.handle,
+            ms=rms.handle,
+            mom=mom.handle,
+            lr=coefficients["lr_t"],
+            rho=coefficients["rho"],
+            momentum=coefficients["momentum"],
+            epsilon=coefficients["epsilon"],
+            grad=grad,
+            use_locking=self._use_locking)
+    else:
+      rms_t = (coefficients["rho"] * rms +
+               coefficients["one_minus_rho"] * math_ops.square(grad))
+      rms_t = state_ops.assign(rms, rms_t, use_locking=self._use_locking)
+      denom_t = rms_t
+      if self.centered:
+        mg = self.get_slot(var, "mg")
+        mg_t = coefficients["rho"] * mg + coefficients["one_minus_rho"] * grad
+        mg_t = state_ops.assign(mg, mg_t, use_locking=self._use_locking)
+        denom_t = rms_t - math_ops.square(mg_t)
+      var_t = var - coefficients["lr_t"] * grad / (
+          math_ops.sqrt(denom_t) + coefficients["epsilon"])
+      return state_ops.assign(var, var_t, use_locking=self._use_locking).op
+
+  def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
+    var_device, var_dtype = var.device, var.dtype.base_dtype
+    coefficients = ((apply_state or {}).get((var_device, var_dtype))
+                    or self._fallback_apply_state(var_device, var_dtype))
+
+    rms = self.get_slot(var, "rms")
+    if self._momentum:
+      mom = self.get_slot(var, "momentum")
+      if self.centered:
+        mg = self.get_slot(var, "mg")
+        return gen_training_ops.ResourceSparseApplyCenteredRMSProp(
+            var=var.handle,
+            mg=mg.handle,
+            ms=rms.handle,
+            mom=mom.handle,
+            lr=coefficients["lr_t"],
+            rho=coefficients["rho"],
+            momentum=coefficients["momentum"],
+            epsilon=coefficients["epsilon"],
+            grad=grad,
+            indices=indices,
+            use_locking=self._use_locking)
+      else:
+        return gen_training_ops.ResourceSparseApplyRMSProp(
+            var=var.handle,
+            ms=rms.handle,
+            mom=mom.handle,
+            lr=coefficients["lr_t"],
+            rho=coefficients["rho"],
+            momentum=coefficients["momentum"],
+            epsilon=coefficients["epsilon"],
+            grad=grad,
+            indices=indices,
+            use_locking=self._use_locking)
+    else:
+      rms_scaled_g_values = (grad * grad) * coefficients["one_minus_rho"]
+      rms_t = state_ops.assign(rms, rms * coefficients["rho"],
+                               use_locking=self._use_locking)
+      with ops.control_dependencies([rms_t]):
+        rms_t = self._resource_scatter_add(rms, indices, rms_scaled_g_values)
+        rms_slice = array_ops.gather(rms_t, indices)
+      denom_slice = rms_slice
+      if self.centered:
+        mg = self.get_slot(var, "mg")
+        mg_scaled_g_values = grad * coefficients["one_minus_rho"]
+        mg_t = state_ops.assign(mg, mg * coefficients["rho"],
+                                use_locking=self._use_locking)
+        with ops.control_dependencies([mg_t]):
+          mg_t = self._resource_scatter_add(mg, indices, mg_scaled_g_values)
+          mg_slice = array_ops.gather(mg_t, indices)
+          denom_slice = rms_slice - math_ops.square(mg_slice)
+      var_update = self._resource_scatter_add(
+          var, indices, coefficients["neg_lr_t"] * grad / (
+              math_ops.sqrt(denom_slice) + coefficients["epsilon"]))
+      if self.centered:
+        return control_flow_ops.group(*[var_update, rms_t, mg_t])
+      return control_flow_ops.group(*[var_update, rms_t])
+
+  def set_weights(self, weights):
+    params = self.weights
+    # Override set_weights for backward compatibility of Keras V1 optimizer
+    # since it does not include iteration at head of the weight list. Set
+    # iteration to 0.
+    if len(params) == len(weights) + 1:
+      weights = [np.array(0)] + weights
+    super(RMSprop, self).set_weights(weights)
+
+  def get_config(self):
+    config = super(RMSprop, self).get_config()
+    config.update({
+        "learning_rate": self._serialize_hyperparameter("learning_rate"),
+        "decay": self._initial_decay,
+        "rho": self._serialize_hyperparameter("rho"),
+        "momentum": self._serialize_hyperparameter("momentum"),
+        "epsilon": self.epsilon,
+        "centered": self.centered,
+    })
+    return config
+
+
+RMSProp = RMSprop

tensorflow/python/keras/optimizer_v2/utils.py

@@ -0,0 +1,157 @@
+# Copyright 2020 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.
+# ==============================================================================
+"""Optimizer utilities."""
+
+from tensorflow.python.distribute import central_storage_strategy
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import reduce_util as ds_reduce_util
+from tensorflow.python.ops import clip_ops
+from tensorflow.python.platform import tf_logging as logging
+
+
+def all_reduce_sum_gradients(grads_and_vars):
+  """Returns all-reduced gradients aggregated via summation.
+
+  Args:
+    grads_and_vars: List of (gradient, variable) pairs.
+
+  Returns:
+    List of (gradient, variable) pairs where gradients have been all-reduced.
+  """
+  grads_and_vars = list(grads_and_vars)
+  filtered_grads_and_vars = filter_empty_gradients(grads_and_vars)
+  if filtered_grads_and_vars:
+    if strategy_supports_no_merge_call():
+      grads = [pair[0] for pair in filtered_grads_and_vars]
+      reduced = distribute_lib.get_strategy().extended._replica_ctx_all_reduce(  # pylint: disable=protected-access
+          ds_reduce_util.ReduceOp.SUM, grads)
+    else:
+      # TODO(b/183257003): Remove this branch
+      reduced = distribute_lib.get_replica_context().merge_call(
+          _all_reduce_sum_fn, args=(filtered_grads_and_vars,))
+  else:
+    reduced = []
+  # Copy 'reduced' but add None gradients back in
+  reduced_with_nones = []
+  reduced_pos = 0
+  for g, v in grads_and_vars:
+    if g is None:
+      reduced_with_nones.append((None, v))
+    else:
+      reduced_with_nones.append((reduced[reduced_pos], v))
+      reduced_pos += 1
+  assert reduced_pos == len(reduced), "Failed to add all gradients"
+  return reduced_with_nones
+
+
+def filter_empty_gradients(grads_and_vars):
+  """Filter out `(grad, var)` pairs that have a gradient equal to `None`."""
+  grads_and_vars = tuple(grads_and_vars)
+  if not grads_and_vars:
+    return grads_and_vars
+
+  filtered = []
+  vars_with_empty_grads = []
+  for grad, var in grads_and_vars:
+    if grad is None:
+      vars_with_empty_grads.append(var)
+    else:
+      filtered.append((grad, var))
+  filtered = tuple(filtered)
+
+  if not filtered:
+    raise ValueError("No gradients provided for any variable: %s." %
+                     ([v.name for _, v in grads_and_vars],))
+  if vars_with_empty_grads:
+    logging.warning(
+        ("Gradients do not exist for variables %s when minimizing the loss."),
+        ([v.name for v in vars_with_empty_grads]))
+  return filtered
+
+
+def make_gradient_clipnorm_fn(clipnorm):
+  """Creates a gradient transformation function for clipping by norm."""
+  if clipnorm is None:
+    return lambda grads_and_vars: grads_and_vars
+
+  def gradient_clipnorm_fn(grads_and_vars):
+
+    if isinstance(distribute_lib.get_strategy(),
+                  (central_storage_strategy.CentralStorageStrategy,
+                   central_storage_strategy.CentralStorageStrategyV1)):
+      raise ValueError(
+          "`clipnorm` is not supported with `CenteralStorageStrategy`")
+
+    clipped_grads_and_vars = [
+        (clip_ops.clip_by_norm(g, clipnorm), v) for g, v in grads_and_vars
+    ]
+    return clipped_grads_and_vars
+
+  return gradient_clipnorm_fn
+
+
+def make_global_gradient_clipnorm_fn(clipnorm):
+  """Creates a gradient transformation function for clipping by norm."""
+  if clipnorm is None:
+    return lambda grads_and_vars: grads_and_vars
+
+  def gradient_clipnorm_fn(grads_and_vars):
+
+    if isinstance(distribute_lib.get_strategy(),
+                  (central_storage_strategy.CentralStorageStrategy,
+                   central_storage_strategy.CentralStorageStrategyV1)):
+      raise ValueError(
+          "`global_clipnorm` is not supported with `CenteralStorageStrategy`")
+
+    grads, variables = zip(*grads_and_vars)
+    clipped_grads, _ = clip_ops.clip_by_global_norm(grads, clipnorm)
+    clipped_grads_and_vars = list(zip(clipped_grads, variables))
+    return clipped_grads_and_vars
+
+  return gradient_clipnorm_fn
+
+
+def make_gradient_clipvalue_fn(clipvalue):
+  """Creates a gradient transformation function for clipping by value."""
+  if clipvalue is None:
+    return lambda grads_and_vars: grads_and_vars
+
+  def gradient_clipvalue_fn(grads_and_vars):
+
+    if isinstance(distribute_lib.get_strategy(),
+                  (central_storage_strategy.CentralStorageStrategy,
+                   central_storage_strategy.CentralStorageStrategyV1)):
+      raise ValueError(
+          "`clipvalue` is not supported with `CenteralStorageStrategy`")
+
+    clipped_grads_and_vars = [(clip_ops.clip_by_value(g, -clipvalue,
+                                                      clipvalue), v)
+                              for g, v in grads_and_vars]
+    return clipped_grads_and_vars
+
+  return gradient_clipvalue_fn
+
+
+def _all_reduce_sum_fn(distribution, grads_and_vars):
+  return distribution.extended.batch_reduce_to(ds_reduce_util.ReduceOp.SUM,
+                                               grads_and_vars)
+
+
+def strategy_supports_no_merge_call():
+  """Returns if the current Strategy can operate in pure replica context."""
+  if not distribute_lib.has_strategy():
+    return True
+  strategy = distribute_lib.get_strategy()
+  return not strategy.extended._use_merge_call()  # pylint: disable=protected-access