text: # 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.
  # ==============================================================================
  """Functional test for GradientDescent."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import tensorflow as tf
  
  from absl.testing import parameterized
  import numpy as np
  from keras import combinations
  from keras.optimizer_v2 import gradient_descent
  from keras.optimizer_v2 import learning_rate_schedule
  
  
  class GradientDescentOptimizerTest(tf.test.TestCase, parameterized.TestCase):
  
    @combinations.generate(combinations.combine(mode=["graph", "eager"]))
    def testBasic(self):
      for dtype in [tf.half, tf.float32, tf.float64]:
        var0 = tf.Variable([1.0, 2.0], dtype=dtype)
        var1 = tf.Variable([3.0, 4.0], dtype=dtype)
        grads0 = tf.constant([0.1, 0.1], dtype=dtype)
        grads1 = tf.constant([0.01, 0.01], dtype=dtype)
        sgd = gradient_descent.SGD(3.0)
        sgd_op = sgd.apply_gradients(zip([grads0, grads1], [var0, var1]))
        self.evaluate(tf.compat.v1.global_variables_initializer())
        # Run 1 step of sgd
        self.evaluate(sgd_op)
        # Validate updated params
        self.assertAllCloseAccordingToType([1.0 - 3.0 * 0.1, 2.0 - 3.0 * 0.1],
                                           self.evaluate(var0))
        self.assertAllCloseAccordingToType([3.0 - 3.0 * 0.01, 4.0 - 3.0 * 0.01],
                                           self.evaluate(var1))
  
    def _test_basic_sgd_with_learning_rate_decay(self, sgd, dtype):
      var0 = tf.Variable([1.0, 2.0], dtype=dtype)
      var1 = tf.Variable([3.0, 4.0], dtype=dtype)
      grads0 = tf.constant([0.1, 0.1], dtype=dtype)
      grads1 = tf.constant([0.01, 0.01], dtype=dtype)
      if not tf.executing_eagerly():
        sgd_op = sgd.apply_gradients(zip([grads0, grads1], [var0, var1]))
      self.evaluate(tf.compat.v1.global_variables_initializer())
      # Run 2 steps of sgd
      if not tf.executing_eagerly():
        self.evaluate(sgd_op)
      else:
        sgd.apply_gradients(zip([grads0, grads1], [var0, var1]))
      # Validate updated params
      self.assertAllCloseAccordingToType([1.0 - 3.0 * 0.1, 2.0 - 3.0 * 0.1],
                                         self.evaluate(var0))
      self.assertAllCloseAccordingToType([3.0 - 3.0 * 0.01, 4.0 - 3.0 * 0.01],
                                         self.evaluate(var1))
  
      if not tf.executing_eagerly():
        self.evaluate(sgd_op)
      else:
        sgd.apply_gradients(zip([grads0, grads1], [var0, var1]))
      # Validate updated params
      self.assertAllCloseAccordingToType(
          [1.0 - 3.0 * 0.1 - 2.0 * 0.1, 2.0 - 3.0 * 0.1 - 2.0 * 0.1],
          self.evaluate(var0))
      self.assertAllCloseAccordingToType(
          [3.0 - 3.0 * 0.01 - 2.0 * 0.01, 4.0 - 3.0 * 0.01 - 2.0 * 0.01],
          self.evaluate(var1))
  
    @combinations.generate(combinations.combine(mode=["graph", "eager"]))
    def testBasicWithLearningRateDecay(self):
      for dtype in [tf.half, tf.float32, tf.float64]:
        learning_rate = 3.0
        decay = 0.5
        sgd = gradient_descent.SGD(learning_rate=learning_rate, decay=decay)
        self._test_basic_sgd_with_learning_rate_decay(sgd, dtype)
  
    @combinations.generate(combinations.combine(mode=["graph", "eager"]))
    def testBasicWithLearningRateInverseTimeDecay(self):
      for dtype in [tf.half, tf.float32, tf.float64]:
        learning_rate = learning_rate_schedule.InverseTimeDecay(
            3.0, decay_steps=1.0, decay_rate=0.5)
        sgd = gradient_descent.SGD(learning_rate=learning_rate)
        self._test_basic_sgd_with_learning_rate_decay(sgd, dtype)
  
    @combinations.generate(combinations.combine(mode=["graph", "eager"]))
    def testBasicWithLearningRateInverseTimeDecaySerializeAndDeserialize(self):
      for dtype in [tf.half, tf.float32, tf.float64]:
        learning_rate = learning_rate_schedule.InverseTimeDecay(
            3.0, decay_steps=1.0, decay_rate=0.5)
        sgd = gradient_descent.SGD(learning_rate=learning_rate)
        sgd = gradient_descent.SGD.from_config(sgd.get_config())
        self._test_basic_sgd_with_learning_rate_decay(sgd, dtype)
  
    @combinations.generate(combinations.combine(mode=["graph", "eager"]))
    def testBasicCallableParams(self):
      for dtype in [tf.half, tf.float32, tf.float64]:
        var0 = tf.Variable([1.0, 2.0], dtype=dtype)
        var1 = tf.Variable([3.0, 4.0], dtype=dtype)
        grads0 = tf.constant([0.1, 0.1], dtype=dtype)
        grads1 = tf.constant([0.01, 0.01], dtype=dtype)
        lr = lambda: 3.0
        sgd = gradient_descent.SGD(lr)
        sgd_op = sgd.apply_gradients(zip([grads0, grads1], [var0, var1]))
        self.evaluate(tf.compat.v1.global_variables_initializer())
        # Run 1 step of sgd
        self.evaluate(sgd_op)
        # Validate updated params
        self.assertAllCloseAccordingToType([1.0 - 3.0 * 0.1, 2.0 - 3.0 * 0.1],
                                           self.evaluate(var0))
        self.assertAllCloseAccordingToType([3.0 - 3.0 * 0.01, 4.0 - 3.0 * 0.01],
                                           self.evaluate(var1))
  
    @combinations.generate(combinations.combine(mode=["graph", "eager"]))
    def testMinimizeResourceVariable(self):
      for dtype in [tf.half, tf.float32, tf.float64]:
        var0 = tf.Variable([[1.0, 2.0]], dtype=dtype)
        var1 = tf.Variable([3.0], dtype=dtype)
        x = tf.constant([[4.0], [5.0]], dtype=dtype)
        loss = lambda: tf.matmul(var0, x) + var1  # pylint: disable=cell-var-from-loop
        sgd = gradient_descent.SGD(1.0)
        sgd_op = sgd.minimize(loss, [var0, var1])
        self.evaluate(tf.compat.v1.global_variables_initializer())
        # Run 1 step of sgd
        self.evaluate(sgd_op)
        # Validate updated params
        self.assertAllCloseAccordingToType([[1.0 - 4.0, 2.0 - 5.0]],
                                           self.evaluate(var0))
        self.assertAllCloseAccordingToType([3.0 - 1.0], self.evaluate(var1))
  
    def testMinimizeSparseResourceVariable(self):
      # TODO(tanzheny, omalleyt): Fix test in eager mode.
      with tf.Graph().as_default():
        for dtype in [tf.half, tf.float32, tf.float64]:
          var0 = tf.Variable([[1.0, 2.0]], dtype=dtype)
          var1 = tf.Variable([3.0], dtype=dtype)
          x = tf.constant([[4.0], [5.0]], dtype=dtype)
  
          def loss():
            pred = tf.matmul(tf.compat.v1.nn.embedding_lookup([var0], [0]), x)  # pylint: disable=cell-var-from-loop
            pred += var1  # pylint: disable=cell-var-from-loop
            return pred * pred
  
          sgd_op = gradient_descent.SGD(1.0).minimize(loss, [var0, var1])
          self.evaluate(tf.compat.v1.global_variables_initializer())
          # Run 1 step of sgd
          self.evaluate(sgd_op)
          # Validate updated params
          np_pred = 1.0 * 4.0 + 2.0 * 5.0 + 3.0
          np_grad = 2 * np_pred
          self.assertAllCloseAccordingToType(
              [[1.0 - np_grad * 4.0, 2.0 - np_grad * 5.0]], self.evaluate(var0))
          self.assertAllCloseAccordingToType([3.0 - np_grad], self.evaluate(var1))
  
    def testTensorLearningRate(self):
      for dtype in [tf.half, tf.float32, tf.float64]:
        var0 = tf.Variable([1.0, 2.0], dtype=dtype)
        var1 = tf.Variable([3.0, 4.0], dtype=dtype)
        grads0 = tf.constant([0.1, 0.1], dtype=dtype)
        grads1 = tf.constant([0.01, 0.01], dtype=dtype)
        lrate = tf.constant(3.0)
        sgd_op = gradient_descent.SGD(lrate).apply_gradients(
            zip([grads0, grads1], [var0, var1]))
        self.evaluate(tf.compat.v1.global_variables_initializer())
        # Run 1 step of sgd
        self.evaluate(sgd_op)
        # Validate updated params
        self.assertAllCloseAccordingToType([1.0 - 3.0 * 0.1, 2.0 - 3.0 * 0.1],
                                           self.evaluate(var0))
        self.assertAllCloseAccordingToType([3.0 - 3.0 * 0.01, 4.0 - 3.0 * 0.01],
                                           self.evaluate(var1))
  
    def testGradWrtRef(self):
      # TODO(tanzheny, omalleyt): Fix test in eager mode.
      with tf.Graph().as_default():
        for dtype in [tf.half, tf.float32, tf.float64]:
          opt = gradient_descent.SGD(3.0)
          values = [1.0, 3.0]
          vars_ = [tf.Variable([v], dtype=dtype) for v in values]
          loss = lambda: vars_[0] + vars_[1]  # pylint: disable=cell-var-from-loop
          grads_and_vars = opt._compute_gradients(loss, vars_)
          self.evaluate(tf.compat.v1.global_variables_initializer())
          for grad, _ in grads_and_vars:
            self.assertAllCloseAccordingToType([1.0], self.evaluate(grad))
  
    def testSparseBasic(self):
      # TODO(tanzheny, omalleyt): Fix test in eager mode.
      with tf.Graph().as_default():
        for dtype in [tf.half, tf.float32, tf.float64]:
          var0 = tf.Variable([[1.0], [2.0]], dtype=dtype)
          var1 = tf.Variable([[3.0], [4.0]], dtype=dtype)
          grads0 = tf.IndexedSlices(
              tf.constant([0.1], shape=[1, 1], dtype=dtype),
              tf.constant([0]), tf.constant([2, 1]))
          grads1 = tf.IndexedSlices(
              tf.constant([0.01], shape=[1, 1], dtype=dtype),
              tf.constant([1]), tf.constant([2, 1]))
          sgd_op = gradient_descent.SGD(3.0).apply_gradients(
              zip([grads0, grads1], [var0, var1]))
          self.evaluate(tf.compat.v1.global_variables_initializer())
          # Run 1 step of sgd
          self.evaluate(sgd_op)
          # Validate updated params
          self.assertAllCloseAccordingToType([[1.0 - 3.0 * 0.1], [2.0]],
                                             self.evaluate(var0))
          self.assertAllCloseAccordingToType([[3.0], [4.0 - 3.0 * 0.01]],
                                             self.evaluate(var1))
  
    def testSparseBasicWithLearningRateDecay(self):
      # TODO(tanzheny, omalleyt): Fix test in eager mode.
      with tf.Graph().as_default():
        for dtype in [tf.half, tf.float32, tf.float64]:
          var0 = tf.Variable([[1.0], [2.0]], dtype=dtype)
          var1 = tf.Variable([[3.0], [4.0]], dtype=dtype)
          grads0 = tf.IndexedSlices(
              tf.constant([0.1], shape=[1, 1], dtype=dtype),
              tf.constant([0]), tf.constant([2, 1]))
          grads1 = tf.IndexedSlices(
              tf.constant([0.01], shape=[1, 1], dtype=dtype),
              tf.constant([1]), tf.constant([2, 1]))
          sgd_op = gradient_descent.SGD(
              3.0, decay=0.5).apply_gradients(
                  zip([grads0, grads1], [var0, var1]))
          self.evaluate(tf.compat.v1.global_variables_initializer())
          # Run 2 steps of sgd
          self.evaluate(sgd_op)
          # Validate updated params
          self.assertAllCloseAccordingToType([[1.0 - 3.0 * 0.1], [2.0]],
                                             self.evaluate(var0))
          self.assertAllCloseAccordingToType([[3.0], [4.0 - 3.0 * 0.01]],
                                             self.evaluate(var1))
  
          self.evaluate(sgd_op)
          # Validate updated params
          self.assertAllCloseAccordingToType(
              [[1.0 - 3.0 * 0.1 - 2.0 * 0.1], [2.0]], self.evaluate(var0))
          self.assertAllCloseAccordingToType(
              [[3.0], [4.0 - 3.0 * 0.01 - 2.0 * 0.01]], self.evaluate(var1))
  
    @combinations.generate(combinations.combine(mode=["eager"]))
    def testCapturingInFunctionWhileExecutingEagerly(self):
      optimizer = gradient_descent.SGD(1.0)
  
      var_holder = {}
      def step():
        if not var_holder:
          var_holder["var"] = tf.Variable(1.0)
        else:
          var_holder["var"].assign(1.0)
  
        with tf.GradientTape() as tape:
          loss = var_holder["var"]**2
        grad = tape.gradient(loss, var_holder["var"])
        optimizer.apply_gradients([(grad, var_holder["var"])])
        return var_holder["var"].read_value()
  
      compiled_step = tf.function(step)
  
      self.assertEqual(float(step()), -1.0)
      self.assertEqual(float(compiled_step()), -1.0)
      # This shouldn't fail; in particular, the learning rate tensor should
      # be an EagerTensor once again, not a graph Tensor.
      self.assertEqual(float(step()), -1.0)
  
    def testConstructSGDWithLR(self):
      opt = gradient_descent.SGD(lr=1.0)
      opt_2 = gradient_descent.SGD(learning_rate=0.1, lr=1.0)
      opt_3 = gradient_descent.SGD(learning_rate=0.1)
      self.assertIsInstance(opt.lr, tf.Variable)
      self.assertIsInstance(opt_2.lr, tf.Variable)
      self.assertIsInstance(opt_3.lr, tf.Variable)
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
      self.assertAllClose(self.evaluate(opt.lr), (1.0))
      self.assertAllClose(self.evaluate(opt_2.lr), (1.0))
      self.assertAllClose(self.evaluate(opt_3.lr), (0.1))
  
  
  class MomentumOptimizerTest(tf.test.TestCase, parameterized.TestCase):
  
    def _update_nesterov_momentum_numpy(self, var, accum, g, lr, momentum):
      accum = accum * momentum - g * lr
      var += (accum * momentum - g * lr)
      return var, accum
  
    @combinations.generate(combinations.combine(mode=["graph", "eager"]))
    def testBasic(self):
      for _, dtype in enumerate([tf.half, tf.float32, tf.float64]):
        var0 = tf.Variable([1.0, 2.0], dtype=dtype, name="var0")
        var1 = tf.Variable([3.0, 4.0], dtype=dtype, name="var1")
        grads0 = tf.constant([0.1, 0.1], dtype=dtype)
        grads1 = tf.constant([0.01, 0.01], dtype=dtype)
        learning_rate = 2.0
        momentum = 0.9
        mom_opt = gradient_descent.SGD(
            learning_rate=learning_rate, momentum=momentum)
        # self.assertFalse(mom_opt._initial_decay)
        mom_update = mom_opt.apply_gradients(
            zip([grads0, grads1], [var0, var1]))
  
        # Check we have slots
        slot0 = mom_opt.get_slot(var0, "momentum")
        self.assertEqual(slot0.shape, var0.shape)
        slot1 = mom_opt.get_slot(var1, "momentum")
        self.assertEqual(slot1.shape, var1.shape)
  
        # Step 1: the momentum accumulators where 0. So we should see a normal
        # update: v -= grad * learning_rate
        self.evaluate(tf.compat.v1.global_variables_initializer())
        self.evaluate(mom_update)
        # Check that the momentum accumulators have been updated.
        self.assertAllCloseAccordingToType(
            np.array([-0.2, -0.2]), self.evaluate(slot0))
        self.assertAllCloseAccordingToType(
            np.array([-0.02, -0.02]), self.evaluate(slot1))
        # Check that the parameters have been updated.
        self.assertAllCloseAccordingToType(
            np.array([1.0 - (0.1 * 2.0), 2.0 - (0.1 * 2.0)]),
            self.evaluate(var0))
        self.assertAllCloseAccordingToType(
            np.array([3.0 - (0.01 * 2.0), 4.0 - (0.01 * 2.0)]),
            self.evaluate(var1))
        # Step 2: the momentum accumulators contain the previous update.
        self.evaluate(mom_update)
        if tf.executing_eagerly():
          mom_opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
        # Check that the momentum accumulators have been updated.
        self.assertAllCloseAccordingToType(
            np.array([(0.9 * (-0.2) - 2.0 * 0.1), (0.9 * (-0.2) - 2.0 * 0.1)]),
            self.evaluate(slot0))
        self.assertAllCloseAccordingToType(
            np.array([(0.9 * (-0.02) - 2.0 * 0.01),
                      (0.9 * (-0.02) - 2.0 * 0.01)]), self.evaluate(slot1))
        # Check that the parameters have been updated.
        self.assertAllCloseAccordingToType(
            np.array([
                1.0 - (0.1 * 2.0) - ((0.9 * 0.1 + 0.1) * 2.0),
                2.0 - (0.1 * 2.0) - ((0.9 * 0.1 + 0.1) * 2.0)
            ]), self.evaluate(var0))
        self.assertAllCloseAccordingToType(
            np.array([
                2.98 - ((0.9 * 0.01 + 0.01) * 2.0),
                3.98 - ((0.9 * 0.01 + 0.01) * 2.0)
            ]), self.evaluate(var1))
  
    def testNesterovMomentum(self):
      # TODO(tanzheny, omalleyt): Fix test in eager mode.
      with tf.Graph().as_default():
        for dtype in [tf.float32, tf.float64]:
          var0 = tf.Variable([1.0, 2.0], dtype=dtype, name="var0")
          var1 = tf.Variable([3.0, 4.0], dtype=dtype, name="var1")
          var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
          var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
          accum0_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
          accum1_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
          loss = lambda: 5 * var0 * var0 + 3 * var1  # pylint: disable=cell-var-from-loop
          mom_op = gradient_descent.SGD(
              learning_rate=2.0, momentum=0.9, nesterov=True)
          opt_op = mom_op.minimize(loss, [var0, var1])
          self.evaluate(tf.compat.v1.global_variables_initializer())
          for _ in range(1, 5):
            self.evaluate(opt_op)
            var0_np, accum0_np = self._update_nesterov_momentum_numpy(
                var0_np, accum0_np, var0_np * 10, 2.0, 0.9)
            var1_np, accum1_np = self._update_nesterov_momentum_numpy(
                var1_np, accum1_np, 3, 2.0, 0.9)
            self.assertAllClose(var0_np, self.evaluate(var0))
            self.assertAllClose(var1_np, self.evaluate(var1))
  
    def testSparseNesterovMomentum(self):
      # TODO(tanzheny, omalleyt): Fix test in eager mode.
      for dtype in [tf.float32, tf.float64]:
        with tf.Graph().as_default(), self.cached_session() as sess:
          var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
          var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
          accum0_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
          accum1_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
          grads = []
          for t in range(1, 5):
            grads.append(var0_np * 10)
            var0_np, accum0_np = self._update_nesterov_momentum_numpy(
                var0_np, accum0_np, var0_np * 10, 2.0, 0.9)
            var1_np, accum1_np = self._update_nesterov_momentum_numpy(
                var1_np, accum1_np, 3, 2.0, 0.9)
          var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
          var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
          accum0_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
          accum1_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
          var0 = tf.Variable(var0_np, dtype=dtype, name="var0")
          var1 = tf.Variable(var1_np, dtype=dtype, name="var1")
          mom_op = gradient_descent.SGD(
              learning_rate=2.0, momentum=0.9, nesterov=True)
          x_feed = tf.compat.v1.placeholder(dtype)
          y_feed = tf.IndexedSlices(x_feed, tf.constant([0, 1]),
                                     tf.constant([2]))
          grads_and_vars = [(y_feed, var0),
                            (tf.constant([3.0, 3.0], dtype=dtype), var1)]
          opt_update = mom_op.apply_gradients(grads_and_vars)
          self.evaluate(tf.compat.v1.global_variables_initializer())
          for t in range(1, 5):
            sess.run(opt_update, feed_dict={x_feed: grads[t - 1]})
            var0_np, accum0_np = self._update_nesterov_momentum_numpy(
                var0_np, accum0_np, var0_np * 10, 2.0, 0.9)
            var1_np, accum1_np = self._update_nesterov_momentum_numpy(
                var1_np, accum1_np, 3, 2.0, 0.9)
            self.assertAllClose(var0_np, self.evaluate(var0))
            self.assertAllClose(var1_np, self.evaluate(var1))
  
    def testMinimizeSparseResourceVariable(self):
      # TODO(tanzheny, omalleyt): Fix test in eager mode.
      with tf.Graph().as_default():
        for dtype in [tf.half, tf.float32, tf.float64]:
          var0 = tf.Variable([[1.0, 2.0]], dtype=dtype)
  
          # pylint: disable=cell-var-from-loop
          def loss():
            x = tf.constant([[4.0], [5.0]], dtype=dtype)
            pred = tf.matmul(tf.compat.v1.nn.embedding_lookup([var0], [0]), x)
            return pred * pred
  
          # pylint: enable=cell-var-from-loop
  
          opt = gradient_descent.SGD(learning_rate=1.0, momentum=0.9)
          sgd_op = opt.minimize(loss, [var0])
          self.evaluate(tf.compat.v1.global_variables_initializer())
          # Run 1 step of sgd
          self.evaluate(sgd_op)
          # Validate updated params
          self.assertAllCloseAccordingToType([[-111, -138]], self.evaluate(var0))
  
    @combinations.generate(combinations.combine(mode=["graph", "eager"]))
    def testMinimizeWith2DIndicesForEmbeddingLookup(self):
      var0 = tf.Variable(tf.ones([2, 2]))
  
      def loss():
        return tf.reduce_sum(tf.compat.v1.nn.embedding_lookup(var0, [[1]]))
  
      opt = gradient_descent.SGD(learning_rate=1.0, momentum=0.9)
      sgd_op = opt.minimize(loss, [var0])
      self.evaluate(tf.compat.v1.global_variables_initializer())
      self.evaluate(sgd_op)
      self.assertAllCloseAccordingToType([[1, 1], [0, 0]], self.evaluate(var0))
  
    def testTensorLearningRateAndMomentum(self):
      # TODO(tanzheny, omalleyt): Fix test in eager mode.
      with tf.Graph().as_default():
        for dtype in [tf.half, tf.float32, tf.float64]:
          var0 = tf.Variable([1.0, 2.0], dtype=dtype)
          var1 = tf.Variable([3.0, 4.0], dtype=dtype)
          grads0 = tf.constant([0.1, 0.1], dtype=dtype)
          grads1 = tf.constant([0.01, 0.01], dtype=dtype)
          mom_opt = gradient_descent.SGD(
              learning_rate=tf.constant(2.0),
              momentum=tf.constant(0.9))
          mom_update = mom_opt.apply_gradients(
              zip([grads0, grads1], [var0, var1]))
          self.evaluate(tf.compat.v1.global_variables_initializer())
          # Check we have slots
          slot0 = mom_opt.get_slot(var0, "momentum")
          self.assertEqual(slot0.shape, var0.shape)
          slot1 = mom_opt.get_slot(var1, "momentum")
          self.assertEqual(slot1.shape, var1.shape)
  
          # Fetch params to validate initial values
          self.assertAllClose([1.0, 2.0], self.evaluate(var0))
          self.assertAllClose([3.0, 4.0], self.evaluate(var1))
          # Step 1: the momentum accumulators where 0. So we should see a normal
          # update: v -= grad * learning_rate
          self.evaluate(mom_update)
          # Check that the momentum accumulators have been updated.
          self.assertAllCloseAccordingToType(
              np.array([-0.2, -0.2]), self.evaluate(slot0))
          self.assertAllCloseAccordingToType(
              np.array([-0.02, -0.02]), self.evaluate(slot1))
          # Check that the parameters have been updated.
          self.assertAllCloseAccordingToType(
              np.array([1.0 - (0.1 * 2.0), 2.0 - (0.1 * 2.0)]),
              self.evaluate(var0))
          self.assertAllCloseAccordingToType(
              np.array([3.0 - (0.01 * 2.0), 4.0 - (0.01 * 2.0)]),
              self.evaluate(var1))
          # Step 2: the momentum accumulators contain the previous update.
          self.evaluate(mom_update)
          # Check that the momentum accumulators have been updated.
          self.assertAllCloseAccordingToType(
              np.array([(0.9 * (-0.2) - 2.0 * 0.1), (0.9 * (-0.2) - 2.0 * 0.1)]),
              self.evaluate(slot0))
          self.assertAllCloseAccordingToType(
              np.array([(0.9 * (-0.02) - 2.0 * 0.01),
                        (0.9 * (-0.02) - 2.0 * 0.01)]), self.evaluate(slot1))
          # Check that the parameters have been updated.
          self.assertAllCloseAccordingToType(
              np.array([
                  1.0 - (0.1 * 2.0) - ((0.9 * 0.1 + 0.1) * 2.0),
                  2.0 - (0.1 * 2.0) - ((0.9 * 0.1 + 0.1) * 2.0)
              ]), self.evaluate(var0))
          self.assertAllCloseAccordingToType(
              np.array([
                  2.98 - ((0.9 * 0.01 + 0.01) * 2.0),
                  3.98 - ((0.9 * 0.01 + 0.01) * 2.0)
              ]), self.evaluate(var1))
  
    def testSparse(self):
      # TODO(tanzheny, omalleyt): Fix test in eager mode.
      with tf.Graph().as_default():
        for dtype in [tf.half, tf.float32, tf.float64]:
          var0 = tf.Variable(tf.zeros([4, 2], dtype=dtype))
          var1 = tf.Variable(tf.constant(1.0, dtype, [4, 2]))
          grads0 = tf.IndexedSlices(
              tf.constant([[.1, .1]], dtype=dtype),
              tf.constant([1]), tf.constant([4, 2]))
          grads1 = tf.IndexedSlices(
              tf.constant([[.01, .01], [.01, .01]], dtype=dtype),
              tf.constant([2, 3]), tf.constant([4, 2]))
          mom_opt = gradient_descent.SGD(learning_rate=2.0, momentum=0.9)
          mom_update = mom_opt.apply_gradients(
              zip([grads0, grads1], [var0, var1]))
          self.evaluate(tf.compat.v1.global_variables_initializer())
  
          # Check we have slots
          slot0 = mom_opt.get_slot(var0, "momentum")
          self.assertEqual(slot0.shape, var0.shape)
          slot1 = mom_opt.get_slot(var1, "momentum")
          self.assertEqual(slot1.shape, var1.shape)
  
          # Fetch params to validate initial values
          self.assertAllClose([0, 0], self.evaluate(var0)[0])
          self.assertAllClose([0, 0], self.evaluate(var0)[1])
          self.assertAllClose([1, 1], self.evaluate(var1)[2])
  
          # Step 1: the momentum accumulators are 0. So we should see a normal
          # update: v -= grad * learning_rate
          self.evaluate(mom_update)
          # Check that the momentum accumulators have been updated.
          self.assertAllCloseAccordingToType(
              np.array([0, 0]),
              self.evaluate(slot0)[0])
          self.assertAllCloseAccordingToType(
              np.array([-2.0 * .1, -2.0 * .1]),
              self.evaluate(slot0)[1])
          self.assertAllCloseAccordingToType(
              np.array([-2.0 * .01, -2.0 * .01]),
              self.evaluate(slot1)[2])
          # Check that the parameters have been updated.
          self.assertAllCloseAccordingToType(
              np.array([0, 0]),
              self.evaluate(var0)[0])
          self.assertAllCloseAccordingToType(
              np.array([-(0.1 * 2.0), -(0.1 * 2.0)]),
              self.evaluate(var0)[1])
          self.assertAllCloseAccordingToType(
              np.array([1.0 - (0.01 * 2.0), 1.0 - (0.01 * 2.0)]),
              self.evaluate(var1)[2])
          # Step 2: the momentum accumulators contain the previous update.
          self.evaluate(mom_update)
          # Check that the momentum accumulators have been updated.
          self.assertAllClose(np.array([0, 0]), self.evaluate(slot0)[0])
          self.assertAllCloseAccordingToType(
              np.array([(0.9 * (-0.2) - 2.0 * 0.1), (0.9 * (-0.2) - 2.0 * 0.1)]),
              self.evaluate(slot0)[1])
          self.assertAllCloseAccordingToType(
              np.array([(0.9 * (-0.02) - 2.0 * 0.01),
                        (0.9 * (-0.02) - 2.0 * 0.01)]),
              self.evaluate(slot1)[2])
          # Check that the parameters have been updated.
          self.assertAllClose(np.array([0, 0]), self.evaluate(var0)[0])
          self.assertAllCloseAccordingToType(
              np.array([
                  -(0.1 * 2.0) - ((0.9 * 0.1 + 0.1) * 2.0),
                  -(0.1 * 2.0) - ((0.9 * 0.1 + 0.1) * 2.0)
              ]),
              self.evaluate(var0)[1])
          self.assertAllCloseAccordingToType(
              np.array([
                  0.98 - ((0.9 * 0.01 + 0.01) * 2.0),
                  0.98 - ((0.9 * 0.01 + 0.01) * 2.0)
              ]),
              self.evaluate(var1)[2])
  
    def testSharing(self):
      # TODO(tanzheny, omalleyt): Fix test in eager mode.
      with tf.Graph().as_default():
        for dtype in [tf.half, tf.float32, tf.float64]:
          var0 = tf.Variable([1.0, 2.0], dtype=dtype)
          var1 = tf.Variable([3.0, 4.0], dtype=dtype)
          grads0 = tf.constant([0.1, 0.1], dtype=dtype)
          grads1 = tf.constant([0.01, 0.01], dtype=dtype)
          mom_opt = gradient_descent.SGD(learning_rate=2.0, momentum=0.9)
          mom_update1 = mom_opt.apply_gradients(
              zip([grads0, grads1], [var0, var1]))
          mom_update2 = mom_opt.apply_gradients(
              zip([grads0, grads1], [var0, var1]))
          self.evaluate(tf.compat.v1.global_variables_initializer())
  
          slot0 = mom_opt.get_slot(var0, "momentum")
          self.assertEqual(slot0.shape, var0.shape)
          slot1 = mom_opt.get_slot(var1, "momentum")
          self.assertEqual(slot1.shape, var1.shape)
  
          # Fetch params to validate initial values
          self.assertAllClose([1.0, 2.0], self.evaluate(var0))
          self.assertAllClose([3.0, 4.0], self.evaluate(var1))
          # Step 1: the momentum accumulators where 0. So we should see a normal
          # update: v -= grad * learning_rate
          self.evaluate(mom_update1)
          # Check that the momentum accumulators have been updated.
          self.assertAllCloseAccordingToType(
              np.array([-0.2, -0.2]), self.evaluate(slot0))
          self.assertAllCloseAccordingToType(
              np.array([-0.02, -0.02]), self.evaluate(slot1))
          # Check that the parameters have been updated.
          self.assertAllCloseAccordingToType(
              np.array([1.0 - (0.1 * 2.0), 2.0 - (0.1 * 2.0)]),
              self.evaluate(var0))
          self.assertAllCloseAccordingToType(
              np.array([3.0 - (0.01 * 2.0), 4.0 - (0.01 * 2.0)]),
              self.evaluate(var1))
          # Step 2: the second momentum accumulators contain the previous update.
          self.evaluate(mom_update2)
          # Check that the momentum accumulators have been updated.
          self.assertAllCloseAccordingToType(
              np.array([(0.9 * (-0.2) - 2.0 * 0.1), (0.9 * (-0.2) - 2.0 * 0.1)]),
              self.evaluate(slot0))
          self.assertAllCloseAccordingToType(
              np.array([(0.9 * (-0.02) - 2.0 * 0.01),
                        (0.9 * (-0.02) - 2.0 * 0.01)]), self.evaluate(slot1))
          # Check that the parameters have been updated.
          self.assertAllCloseAccordingToType(
              np.array([
                  1.0 - (0.1 * 2.0) - ((0.9 * 0.1 + 0.1) * 2.0),
                  2.0 - (0.1 * 2.0) - ((0.9 * 0.1 + 0.1) * 2.0)
              ]), self.evaluate(var0))
          self.assertAllCloseAccordingToType(
              np.array([
                  2.98 - ((0.9 * 0.01 + 0.01) * 2.0),
                  3.98 - ((0.9 * 0.01 + 0.01) * 2.0)
              ]), self.evaluate(var1))
  
    @combinations.generate(combinations.combine(mode=["graph", "eager"]))
    def testConfig(self):
      opt = gradient_descent.SGD(learning_rate=1.0, momentum=0.9, nesterov=True)
      config = opt.get_config()
      opt2 = gradient_descent.SGD.from_config(config)
      lr = opt.lr
      lr2 = opt2.lr
      self.evaluate(tf.compat.v1.global_variables_initializer())
      self.assertAllClose(self.evaluate(lr), self.evaluate(lr2))
      self.assertAllClose(
          self.evaluate(opt._get_hyper("momentum")),
          self.evaluate(opt2._get_hyper("momentum")))
      self.assertAllClose(
          self.evaluate(opt._get_hyper("decay")),
          self.evaluate(opt2._get_hyper("decay")))
      var0 = tf.Variable([[1.0], [2.0]], dtype=tf.float32)
      loss = lambda: 3 * var0
      # learning rate variable created when calling minimize.
      opt.minimize(loss, [var0])
      self.evaluate(tf.compat.v1.global_variables_initializer())
      config = opt.get_config()
      opt3 = gradient_descent.SGD.from_config(config)
      lr3 = opt3.lr
      self.evaluate(tf.compat.v1.global_variables_initializer())
      self.assertAllClose(self.evaluate(lr), self.evaluate(lr3))
      self.assertAllClose(
          self.evaluate(opt._get_hyper("momentum")),
          self.evaluate(opt3._get_hyper("momentum")))
      self.assertAllClose(
          self.evaluate(opt._get_hyper("decay")),
          self.evaluate(opt3._get_hyper("decay")))
      self.assertTrue(opt3.nesterov)
  
    def testNesterovWithoutMomentum(self):
      with self.assertRaisesRegex(ValueError, "must be between"):
        gradient_descent.SGD(learning_rate=1.0, momentum=2.0)
  
    def testConstructMomentumWithLR(self):
      opt = gradient_descent.SGD(lr=1.0, momentum=0.9)
      opt_2 = gradient_descent.SGD(learning_rate=0.1, momentum=0.9, lr=1.0)
      opt_3 = gradient_descent.SGD(learning_rate=0.1, momentum=0.9)
      self.assertIsInstance(opt.lr, tf.Variable)
      self.assertIsInstance(opt_2.lr, tf.Variable)
      self.assertIsInstance(opt_3.lr, tf.Variable)
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
      self.assertAllClose(self.evaluate(opt.lr), (1.0))
      self.assertAllClose(self.evaluate(opt_2.lr), (1.0))
      self.assertAllClose(self.evaluate(opt_3.lr), (0.1))
  
    @combinations.generate(combinations.combine(mode=["eager"]))
    def testMinimizeLossTensor(self):
      for dtype in [tf.half, tf.float32, tf.float64]:
        var0 = tf.Variable([[1.0, 2.0]], dtype=dtype)
        var1 = tf.Variable([3.0], dtype=dtype)
        x = tf.constant([[4.0], [5.0]], dtype=dtype)
  
        tape = tf.GradientTape()
        with tape:
          loss = tf.matmul(var0, x) + var1
        sgd = gradient_descent.SGD(1.0)
        with self.assertRaisesRegex(ValueError, "`tape` is required"):
          sgd.minimize(loss, [var0, var1])
        sgd.minimize(loss, [var0, var1], tape=tape)
  
        self.assertAllCloseAccordingToType([[1.0 - 4.0, 2.0 - 5.0]],
                                           self.evaluate(var0))
        self.assertAllCloseAccordingToType([3.0 - 1.0], self.evaluate(var1))
  
  
  if __name__ == "__main__":
    tf.test.main()
  