chore: import upstream snapshot with attribution

tensorflow/python/training/adadelta.py

@@ -0,0 +1,198 @@
+# 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.
+# ==============================================================================
+
+"""Adadelta for TensorFlow."""
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import gen_training_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.training import optimizer
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(v1=["train.AdadeltaOptimizer"])
+class AdadeltaOptimizer(optimizer.Optimizer):
+  """Optimizer that implements the Adadelta algorithm.
+
+  References:
+    ADADELTA - An Adaptive Learning Rate Method:
+      [Zeiler, 2012](http://arxiv.org/abs/1212.5701)
+      ([pdf](http://arxiv.org/pdf/1212.5701v1.pdf))
+
+  @compatibility(TF2)
+  tf.compat.v1.train.AdadeltaOptimizer is compatible with eager mode and
+  `tf.function`.
+  When eager execution is enabled, `learning_rate`, `rho`,
+  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.
+
+  To switch to native TF2 style, use [`tf.keras.optimizers.Adadelta`]
+  (https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adadelta)
+  instead. Please notice that due to the implementation differences,
+  `tf.keras.optimizers.Adadelta` and
+  `tf.compat.v1.train.AdadeltaOptimizer` may have slight differences in
+  floating point numerics even though the formula used for the variable
+  updates still matches.
+
+  #### Structural mapping to native TF2
+
+  Before:
+
+  ```python
+  optimizer = tf.compat.v1.train.AdadeltaOptimizer(
+    learning_rate=learning_rate,
+    rho=rho,
+    epsilon=epsilon)
+  ```
+
+  After:
+
+  ```python
+  optimizer = tf.keras.optimizers.Adadelta(
+    learning_rate=learning_rate,
+    rho=rho,
+    epsilon=epsilon)
+  ```
+
+  #### How to map arguments
+  | TF1 Arg Name       | TF2 Arg Name   | Note                             |
+  | ------------------ | -------------  | -------------------------------  |
+  | `learning_rate`    | `learning_rate`| Be careful of setting           |
+  : : : learning_rate tensor value computed from the global step.          :
+  : : : In TF1 this was usually meant to imply a dynamic learning rate and :
+  : : : would recompute in each step. In TF2 (eager + function) it will    :
+  : : : treat it as a scalar value that only gets computed once instead of :
+  : : : a symbolic placeholder to be computed each time.                   :
+  | `rho`              | `rho`          | -                                |
+  | `epsilon`          | `epsilon`      | Default value is 1e-08 in TF1,   |
+  :                    :                : but 1e-07 in TF2.                :
+  | `use_locking`      | -              | Not applicable in TF2.           |
+
+  #### Before & after usage example
+  Before:
+
+  ```python
+  x = tf.Variable([1,2,3], dtype=tf.float32)
+  grad = tf.constant([0.1, 0.2, 0.3])
+  optimizer = tf.compat.v1.train.AdadeltaOptimizer(learning_rate=0.001)
+  optimizer.apply_gradients(zip([grad], [x]))
+  ```
+
+  After:
+
+  ```python
+  x = tf.Variable([1,2,3], dtype=tf.float32)
+  grad = tf.constant([0.1, 0.2, 0.3])
+  optimizer = tf.keras.optimizers.Adadelta(learning_rate=0.001)
+  optimizer.apply_gradients(zip([grad], [x]))
+  ```
+
+  @end_compatibility
+  """
+
+  def __init__(self, learning_rate=0.001, rho=0.95, epsilon=1e-8,
+               use_locking=False, name="Adadelta"):
+    """Construct a new Adadelta optimizer.
+
+    Args:
+      learning_rate: A `Tensor` or a floating point value. The learning rate.
+        To match the exact form in the original paper use 1.0.
+      rho: A `Tensor` or a floating point value. The decay rate.
+      epsilon: A `Tensor` or a floating point value.  A constant epsilon used
+               to better conditioning the grad update.
+      use_locking: If `True` use locks for update operations.
+      name: Optional name prefix for the operations created when applying
+        gradients.  Defaults to "Adadelta".
+
+
+    """
+    super(AdadeltaOptimizer, self).__init__(use_locking, name)
+    self._lr = learning_rate
+    self._rho = rho
+    self._epsilon = epsilon
+
+    # Tensor versions of the constructor arguments, created in _prepare().
+    self._lr_t = None
+    self._rho_t = None
+    self._epsilon_t = None
+
+  def _create_slots(self, var_list):
+    for v in var_list:
+      self._zeros_slot(v, "accum", self._name)
+      self._zeros_slot(v, "accum_update", self._name)
+
+  def _prepare(self):
+    lr = self._call_if_callable(self._lr)
+    rho = self._call_if_callable(self._rho)
+    epsilon = self._call_if_callable(self._epsilon)
+
+    self._lr_t = ops.convert_to_tensor(lr, name="lr")
+    self._rho_t = ops.convert_to_tensor(rho, name="rho")
+    self._epsilon_t = ops.convert_to_tensor(epsilon, name="epsilon")
+
+  def _apply_dense(self, grad, var):
+    accum = self.get_slot(var, "accum")
+    accum_update = self.get_slot(var, "accum_update")
+    return gen_training_ops.apply_adadelta(
+        var,
+        accum,
+        accum_update,
+        math_ops.cast(self._lr_t, var.dtype.base_dtype),
+        math_ops.cast(self._rho_t, var.dtype.base_dtype),
+        math_ops.cast(self._epsilon_t, var.dtype.base_dtype),
+        grad,
+        use_locking=self._use_locking)
+
+  def _resource_apply_dense(self, grad, var):
+    accum = self.get_slot(var, "accum")
+    accum_update = self.get_slot(var, "accum_update")
+    return gen_training_ops.resource_apply_adadelta(
+        var.handle,
+        accum.handle,
+        accum_update.handle,
+        math_ops.cast(self._lr_t, grad.dtype.base_dtype),
+        math_ops.cast(self._rho_t, grad.dtype.base_dtype),
+        math_ops.cast(self._epsilon_t, grad.dtype.base_dtype),
+        grad,
+        use_locking=self._use_locking)
+
+  def _apply_sparse(self, grad, var):
+    accum = self.get_slot(var, "accum")
+    accum_update = self.get_slot(var, "accum_update")
+    return gen_training_ops.sparse_apply_adadelta(
+        var,
+        accum,
+        accum_update,
+        math_ops.cast(self._lr_t, var.dtype.base_dtype),
+        math_ops.cast(self._rho_t, var.dtype.base_dtype),
+        math_ops.cast(self._epsilon_t, var.dtype.base_dtype),
+        grad.values,
+        grad.indices,
+        use_locking=self._use_locking)
+
+  def _resource_apply_sparse(self, grad, var, indices):
+    accum = self.get_slot(var, "accum")
+    accum_update = self.get_slot(var, "accum_update")
+    return gen_training_ops.resource_sparse_apply_adadelta(
+        var.handle,
+        accum.handle,
+        accum_update.handle,
+        math_ops.cast(self._lr_t, grad.dtype),
+        math_ops.cast(self._rho_t, grad.dtype),
+        math_ops.cast(self._epsilon_t, grad.dtype),
+        grad,
+        indices,
+        use_locking=self._use_locking)

tensorflow/python/training/adadelta_test.py

@@ -0,0 +1,185 @@
+# 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.
+# ==============================================================================
+"""Tests for Adadelta Optimizer."""
+
+import numpy as np
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import embedding_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import adadelta
+
+
+class AdadeltaOptimizerTest(test.TestCase):
+
+  def doTestBasic(self, use_resource=False, use_callable_params=False):
+    num_updates = 4  # number of ADADELTA steps to perform
+    for dtype in [dtypes.half, dtypes.float32]:
+      for grad in [0.2, 0.1, 0.01]:
+        for lr in [1.0, 0.5, 0.1]:
+          var0_init = [1.0, 2.0]
+          var1_init = [3.0, 4.0]
+          if use_resource:
+            var0 = resource_variable_ops.ResourceVariable(
+                var0_init, dtype=dtype)
+            var1 = resource_variable_ops.ResourceVariable(
+                var1_init, dtype=dtype)
+          else:
+            var0 = variables.Variable(var0_init, dtype=dtype)
+            var1 = variables.Variable(var1_init, dtype=dtype)
+
+          grads = constant_op.constant([grad, grad], dtype=dtype)
+
+          accum = 0.0
+          accum_update = 0.0
+
+          # ADADELTA gradient optimizer
+          rho = 0.95
+          epsilon = 1e-8
+          if use_callable_params:
+            adadelta_opt = adadelta.AdadeltaOptimizer(
+                learning_rate=lambda: lr,  # pylint: disable=cell-var-from-loop
+                rho=lambda: rho,  # pylint: disable=cell-var-from-loop
+                epsilon=lambda: epsilon)  # pylint: disable=cell-var-from-loop
+          else:
+            adadelta_opt = adadelta.AdadeltaOptimizer(
+                learning_rate=lr, rho=rho, epsilon=epsilon)
+          if not context.executing_eagerly():
+            adadelta_update = adadelta_opt.apply_gradients(
+                zip([grads, grads], [var0, var1]))
+            self.evaluate(variables.global_variables_initializer())
+
+            # TODO(lxuechen): This is hard to test in eager mode,
+            # since the optimizer is not fully initialized until the first
+            # call to `apply_gradients`
+            opt_vars = adadelta_opt.variables()
+            self.assertStartsWith(opt_vars[0].name, var0._shared_name)
+            self.assertStartsWith(opt_vars[1].name, var0._shared_name)
+            self.assertStartsWith(opt_vars[2].name, var1._shared_name)
+            self.assertStartsWith(opt_vars[3].name, var1._shared_name)
+            self.assertEqual(4, len(opt_vars))
+            # Assign slots
+            slot = [None] * 2
+            slot_update = [None] * 2
+            self.assertEqual(["accum", "accum_update"],
+                             adadelta_opt.get_slot_names())
+            slot[0] = adadelta_opt.get_slot(var0, "accum")
+            self.assertEqual(slot[0].get_shape(), var0.get_shape())
+            self.assertFalse(slot[0] in variables.trainable_variables())
+
+            slot_update[0] = adadelta_opt.get_slot(var0, "accum_update")
+            self.assertEqual(slot_update[0].get_shape(), var0.get_shape())
+            self.assertFalse(slot_update[0] in variables.trainable_variables())
+
+            slot[1] = adadelta_opt.get_slot(var1, "accum")
+            self.assertEqual(slot[1].get_shape(), var1.get_shape())
+            self.assertFalse(slot[1] in variables.trainable_variables())
+
+            slot_update[1] = adadelta_opt.get_slot(var1, "accum_update")
+            self.assertEqual(slot_update[1].get_shape(), var1.get_shape())
+            self.assertFalse(slot_update[1] in variables.trainable_variables())
+
+          # Fetch params to validate initial values
+          self.assertAllClose(var0_init, self.evaluate(var0))
+          self.assertAllClose(var1_init, self.evaluate(var1))
+
+          update = [None] * num_updates
+          tot_update = 0
+          for step in range(num_updates):
+            # Run adadelta update for comparison
+            if not context.executing_eagerly():
+              self.evaluate(adadelta_update)
+            else:
+              adadelta_opt.apply_gradients(zip([grads, grads], [var0, var1]))
+
+            # Perform initial update without previous accum values
+            accum = accum * rho + (grad**2) * (1 - rho)
+            update[step] = (
+                np.sqrt(accum_update + epsilon) *
+                (1. / np.sqrt(accum + epsilon)) * grad)
+            accum_update = (
+                accum_update * rho + (update[step]**2) * (1.0 - rho))
+            tot_update += update[step] * lr
+
+            if not context.executing_eagerly():
+              # Check that the accumulators have been updated
+              # TODO(lxuechen): This is hard to test in eager mode
+              for slot_idx in range(2):
+                self.assertAllCloseAccordingToType(
+                    np.array([accum, accum], dtype=dtype.as_numpy_dtype()),
+                    self.evaluate(slot[slot_idx]),
+                    rtol=1e-5)
+
+                self.assertAllCloseAccordingToType(
+                    np.array(
+                        [accum_update, accum_update],
+                        dtype=dtype.as_numpy_dtype()),
+                    self.evaluate(slot_update[slot_idx]),
+                    rtol=1e-5)
+
+              # Check that the parameters have been updated
+              self.assertAllCloseAccordingToType(
+                  np.array(
+                      [var0_init[0] - tot_update, var0_init[1] - tot_update],
+                      dtype=dtype.as_numpy_dtype()),
+                  self.evaluate(var0),
+                  rtol=1e-5)
+
+              self.assertAllCloseAccordingToType(
+                  np.array(
+                      [var1_init[0] - tot_update, var1_init[1] - tot_update],
+                      dtype=dtype.as_numpy_dtype()),
+                  self.evaluate(var1),
+                  rtol=1e-5)
+
+  def testBasic(self):
+    with self.cached_session():
+      self.doTestBasic(use_resource=False)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testResourceBasic(self):
+    self.doTestBasic(use_resource=True)
+
+  def testBasicCallableParams(self):
+    with context.eager_mode():
+      self.doTestBasic(use_resource=True, use_callable_params=True)
+
+  @test_util.run_deprecated_v1
+  def testMinimizeSparseResourceVariable(self):
+    for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+      with self.cached_session():
+        var0 = resource_variable_ops.ResourceVariable([[1.0, 2.0]], dtype=dtype)
+        x = constant_op.constant([[4.0], [5.0]], dtype=dtype)
+        pred = math_ops.matmul(embedding_ops.embedding_lookup([var0], [0]), x)
+        loss = pred * pred
+        sgd_op = adadelta.AdadeltaOptimizer(
+            1.0, 1.0, 1.0).minimize(loss)
+        self.evaluate(variables.global_variables_initializer())
+        # Fetch params to validate initial values
+        self.assertAllCloseAccordingToType([[1.0, 2.0]], self.evaluate(var0))
+        # Run 1 step of sgd
+        sgd_op.run()
+        # Validate updated params
+        self.assertAllCloseAccordingToType([[-111, -138]], self.evaluate(var0))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/adagrad.py

@@ -0,0 +1,195 @@
+# 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.
+# ==============================================================================
+
+"""Adagrad for TensorFlow."""
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_array_ops
+from tensorflow.python.ops import gen_training_ops
+from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.training import optimizer
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(v1=["train.AdagradOptimizer"])
+class AdagradOptimizer(optimizer.Optimizer):
+  """Optimizer that implements the Adagrad algorithm.
+
+  References:
+    Adaptive Subgradient Methods for Online Learning and Stochastic Optimization
+      :[Duchi et al., 2011](http://jmlr.org/papers/v12/duchi11a.html)
+      ([pdf](http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf))
+
+  @compatibility(TF2)
+  tf.compat.v1.train.AdagradOptimizer is compatible with eager mode and
+  `tf.function`.
+  When eager execution is enabled, `learning_rate`,
+  `initial_accumulator_value`, 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.
+
+  To switch to native TF2 style, use [`tf.keras.optimizers.Adagrad`]
+  (https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adagrad)
+  instead. Please notice that due to the implementation differences,
+  `tf.keras.optimizers.Adagrad` and
+  `tf.compat.v1.train.AdagradOptimizer` may have slight differences in
+  floating point numerics even though the formula used for the variable
+  updates still matches.
+
+  #### Structural mapping to native TF2
+
+  Before:
+
+  ```python
+  optimizer = tf.compat.v1.train.AdagradOptimizer(
+    learning_rate=learning_rate,
+    initial_accumulator_value=initial_accumulator_value)
+  ```
+
+  After:
+
+  ```python
+  optimizer = tf.keras.optimizers.Adagrad(
+    learning_rate=learning_rate,
+    initial_accumulator_value=initial_accumulator_value,
+    epsilon=1e-07)
+  ```
+
+  #### How to map arguments
+  | TF1 Arg Name       | TF2 Arg Name   | Note                             |
+  | ------------------ | -------------  | -------------------------------  |
+  | `learning_rate`    | `learning_rate` | Be careful of setting           |
+  : : : learning_rate tensor value computed from the global step.          :
+  : : : In TF1 this was usually meant to imply a dynamic learning rate and :
+  : : : would recompute in each step. In TF2 (eager + function) it will    :
+  : : : treat it as a scalar value that only gets computed once instead of :
+  : : : a symbolic placeholder to be computed each time.                   :
+  | `initial_accumulator_value` | `initial_accumulator_value` | The        |
+  : : : argument can be value of zero in TF2, which is not accepted in TF1.|
+  | - | `epsilon`      | `epsilon` is become configurable in TF2. The      |
+  : : : defualt value is changed from 1e-8 to 1e-7                         :
+  | `use_locking`      | -             | Not applicable in TF2.            |
+
+  #### Before & after usage example
+  Before:
+
+  ```python
+  x = tf.Variable([1,2,3], dtype=tf.float32)
+  grad = tf.constant([0.1, 0.2, 0.3])
+  optimizer = tf.compat.v1.train.AdagradOptimizer(learning_rate=0.001)
+  optimizer.apply_gradients(zip([grad], [x]))
+  ```
+
+  After:
+
+  ```python
+  x = tf.Variable([1,2,3], dtype=tf.float32)
+  grad = tf.constant([0.1, 0.2, 0.3])
+  optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.001)
+  optimizer.apply_gradients(zip([grad], [x]))
+  ```
+
+  @end_compatibility
+  """
+
+  def __init__(self, learning_rate, initial_accumulator_value=0.1,
+               use_locking=False, name="Adagrad"):
+    """Construct a new Adagrad optimizer.
+
+    Args:
+      learning_rate: A `Tensor` or a floating point value.  The learning rate.
+      initial_accumulator_value: A floating point value.
+        Starting value for the accumulators, must be positive.
+      use_locking: If `True` use locks for update operations.
+      name: Optional name prefix for the operations created when applying
+        gradients.  Defaults to "Adagrad".
+
+    Raises:
+      ValueError: If the `initial_accumulator_value` is invalid.
+
+    """
+    if initial_accumulator_value <= 0.0:
+      raise ValueError("initial_accumulator_value must be positive: %s" %
+                       initial_accumulator_value)
+    super(AdagradOptimizer, self).__init__(use_locking, name)
+    self._learning_rate = learning_rate
+    self._initial_accumulator_value = initial_accumulator_value
+    # Created in Initialize.
+    self._learning_rate_tensor = None
+
+  def _create_slots(self, var_list):
+    for v in var_list:
+      dtype = v.dtype.base_dtype
+      if v.get_shape().is_fully_defined():
+        init = init_ops.constant_initializer(self._initial_accumulator_value,
+                                             dtype=dtype)
+      else:
+        init = self._init_constant_op(v, dtype)
+      self._get_or_make_slot_with_initializer(v, init, v.get_shape(), dtype,
+                                              "accumulator", self._name)
+
+  def _init_constant_op(self, v, dtype):
+    def init():
+      # Use a Tensor instead of initializer if variable does not have
+      # static shape.
+      init_constant = gen_array_ops.fill(array_ops.shape(v),
+                                         self._initial_accumulator_value)
+      return math_ops.cast(init_constant, dtype)
+    return init
+
+  def _prepare(self):
+    learning_rate = self._call_if_callable(self._learning_rate)
+    self._learning_rate_tensor = ops.convert_to_tensor(
+        learning_rate, name="learning_rate")
+
+  def _apply_dense(self, grad, var):
+    acc = self.get_slot(var, "accumulator")
+    return gen_training_ops.apply_adagrad(
+        var,
+        acc,
+        math_ops.cast(self._learning_rate_tensor, var.dtype.base_dtype),
+        grad,
+        use_locking=self._use_locking)
+
+  def _resource_apply_dense(self, grad, var):
+    acc = self.get_slot(var, "accumulator")
+    return gen_training_ops.resource_apply_adagrad(
+        var.handle,
+        acc.handle,
+        math_ops.cast(self._learning_rate_tensor, grad.dtype.base_dtype),
+        grad,
+        use_locking=self._use_locking)
+
+  def _apply_sparse(self, grad, var):
+    acc = self.get_slot(var, "accumulator")
+    return gen_training_ops.sparse_apply_adagrad(
+        var,
+        acc,
+        math_ops.cast(self._learning_rate_tensor, var.dtype.base_dtype),
+        grad.values,
+        grad.indices,
+        use_locking=self._use_locking)
+
+  def _resource_apply_sparse(self, grad, var, indices):
+    acc = self.get_slot(var, "accumulator")
+    return gen_training_ops.resource_sparse_apply_adagrad(
+        var.handle,
+        acc.handle,
+        math_ops.cast(self._learning_rate_tensor, grad.dtype),
+        grad,
+        indices,
+        use_locking=self._use_locking)

tensorflow/python/training/adagrad_da.py

@@ -0,0 +1,171 @@
+# Copyright 2016 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 Dual Averaging for TensorFlow."""
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_training_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.training import optimizer
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(v1=["train.AdagradDAOptimizer"])
+class AdagradDAOptimizer(optimizer.Optimizer):
+  """Adagrad Dual Averaging algorithm for sparse linear models.
+
+  This optimizer takes care of regularization of unseen features in a mini batch
+  by updating them when they are seen with a closed form update rule that is
+  equivalent to having updated them on every mini-batch.
+
+  AdagradDA is typically used when there is a need for large sparsity in the
+  trained model. This optimizer only guarantees sparsity for linear models. Be
+  careful when using AdagradDA for deep networks as it will require careful
+  initialization of the gradient accumulators for it to train.
+
+  References:
+    Adaptive Subgradient Methods for Online Learning and Stochastic Optimization
+      :[Duchi et al., 2011](http://jmlr.org/papers/v12/duchi11a.html)
+      ([pdf](http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf))
+  """
+
+  def __init__(self,
+               learning_rate,
+               global_step,
+               initial_gradient_squared_accumulator_value=0.1,
+               l1_regularization_strength=0.0,
+               l2_regularization_strength=0.0,
+               use_locking=False,
+               name="AdagradDA"):
+    """Construct a new AdagradDA optimizer.
+
+    Args:
+      learning_rate: A `Tensor` or a floating point value.  The learning rate.
+      global_step: A `Tensor` containing the current training step number.
+      initial_gradient_squared_accumulator_value: A floating point value.
+        Starting value for the accumulators, must be positive.
+      l1_regularization_strength: A float value, must be greater than or
+        equal to zero.
+      l2_regularization_strength: A float value, must be greater than or
+        equal to zero.
+      use_locking: If `True` use locks for update operations.
+      name: Optional name prefix for the operations created when applying
+        gradients.  Defaults to "AdagradDA".
+
+    Raises:
+      ValueError: If the `initial_gradient_squared_accumulator_value` is
+      invalid.
+    """
+    if initial_gradient_squared_accumulator_value <= 0.0:
+      raise ValueError("initial_gradient_squared_accumulator_value must be "
+                       "positive: %s" %
+                       initial_gradient_squared_accumulator_value)
+    super(AdagradDAOptimizer, self).__init__(use_locking, name)
+    self._learning_rate = learning_rate
+    self._initial_gradient_squared_accumulator_value = (
+        initial_gradient_squared_accumulator_value)
+    # Created in Initialize.
+    self._learning_rate_tensor = None
+    self._l1_regularization_strength = l1_regularization_strength
+    self._l2_regularization_strength = l2_regularization_strength
+    self._global_step = global_step
+    self._global_step_on_worker = None
+
+  def _create_slots(self, var_list):
+    for v in var_list:
+      with ops.colocate_with(v):
+        g_val = constant_op.constant(
+            0.0, shape=v.get_shape(), dtype=v.dtype.base_dtype)
+        gg_val = constant_op.constant(
+            self._initial_gradient_squared_accumulator_value,
+            shape=v.get_shape(),
+            dtype=v.dtype.base_dtype)
+      self._get_or_make_slot(v, g_val, "gradient_accumulator", self._name)
+      self._get_or_make_slot(v, gg_val, "gradient_squared_accumulator",
+                             self._name)
+
+  def _prepare(self):
+    self._learning_rate_tensor = ops.convert_to_tensor(
+        self._learning_rate, name="learning_rate")
+    # Performance optimization so that worker creates a copy of the global step
+    # to avoid overloading the parameter server holding the global step.
+    with ops.colocate_with(self._learning_rate_tensor):
+      self._global_step_on_worker = array_ops.identity(self._global_step) + 1
+
+  def _apply_dense(self, grad, var):
+    g_acc = self.get_slot(var, "gradient_accumulator")
+    gg_acc = self.get_slot(var, "gradient_squared_accumulator")
+    with ops.device(var.device):
+      global_step = array_ops.identity(self._global_step_on_worker)
+    return gen_training_ops.apply_adagrad_da(
+        var,
+        g_acc,
+        gg_acc,
+        grad,
+        math_ops.cast(self._learning_rate_tensor, var.dtype.base_dtype),
+        math_ops.cast(self._l1_regularization_strength, var.dtype.base_dtype),
+        math_ops.cast(self._l2_regularization_strength, var.dtype.base_dtype),
+        global_step,
+        use_locking=self._use_locking)
+
+  def _resource_apply_dense(self, grad, var):
+    g_acc = self.get_slot(var, "gradient_accumulator")
+    gg_acc = self.get_slot(var, "gradient_squared_accumulator")
+    with ops.device(var.device):
+      global_step = array_ops.identity(self._global_step_on_worker)
+    return gen_training_ops.resource_apply_adagrad_da(
+        var.handle,
+        g_acc.handle,
+        gg_acc.handle,
+        grad,
+        math_ops.cast(self._learning_rate_tensor, grad.dtype.base_dtype),
+        math_ops.cast(self._l1_regularization_strength, grad.dtype.base_dtype),
+        math_ops.cast(self._l2_regularization_strength, grad.dtype.base_dtype),
+        global_step,
+        use_locking=self._use_locking)
+
+  def _apply_sparse(self, grad, var):
+    g_acc = self.get_slot(var, "gradient_accumulator")
+    gg_acc = self.get_slot(var, "gradient_squared_accumulator")
+    with ops.device(var.device):
+      global_step = array_ops.identity(self._global_step_on_worker)
+    return gen_training_ops.sparse_apply_adagrad_da(
+        var,
+        g_acc,
+        gg_acc,
+        grad.values,
+        grad.indices,
+        math_ops.cast(self._learning_rate_tensor, var.dtype.base_dtype),
+        math_ops.cast(self._l1_regularization_strength, var.dtype.base_dtype),
+        math_ops.cast(self._l2_regularization_strength, var.dtype.base_dtype),
+        global_step,
+        use_locking=self._use_locking)
+
+  def _resource_apply_sparse(self, grad, var, indices):
+    g_acc = self.get_slot(var, "gradient_accumulator")
+    gg_acc = self.get_slot(var, "gradient_squared_accumulator")
+    with ops.device(var.device):
+      global_step = array_ops.identity(self._global_step_on_worker)
+    return gen_training_ops.resource_sparse_apply_adagrad_da(
+        var.handle,
+        g_acc.handle,
+        gg_acc.handle,
+        grad,
+        indices,
+        math_ops.cast(self._learning_rate_tensor, grad.dtype),
+        math_ops.cast(self._l1_regularization_strength, grad.dtype),
+        math_ops.cast(self._l2_regularization_strength, grad.dtype),
+        global_step,
+        use_locking=self._use_locking)

tensorflow/python/training/adagrad_da_test.py

@@ -0,0 +1,200 @@
+# Copyright 2016 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 tests for AdagradDA operations."""
+
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import embedding_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import adagrad_da
+
+
+class AdagradDAOptimizerTest(test.TestCase):
+
+  def doTestAdagradDAwithoutRegularizationBasic1(self, use_resource=False):
+    for dtype in [dtypes.float64, dtypes.float32]:
+      with ops.Graph().as_default(), self.cached_session():
+        global_step = variables.Variable(0, dtype=dtypes.int64)
+        if use_resource:
+          var0 = resource_variable_ops.ResourceVariable([0.0, 0.0], dtype=dtype)
+          var1 = resource_variable_ops.ResourceVariable([0.0, 0.0], dtype=dtype)
+        else:
+          var0 = variables.Variable([0.0, 0.0], dtype=dtype)
+          var1 = variables.Variable([0.0, 0.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+        opt = adagrad_da.AdagradDAOptimizer(
+            3.0,
+            global_step,
+            initial_gradient_squared_accumulator_value=0.1,
+            l1_regularization_strength=0.0,
+            l2_regularization_strength=0.0)
+        update = opt.apply_gradients(
+            zip([grads0, grads1], [var0, var1]), global_step=global_step)
+        self.evaluate(variables.global_variables_initializer())
+
+        v0_val, v1_val = self.evaluate([var0, var1])
+        self.assertAllClose([0.0, 0.0], v0_val)
+        self.assertAllClose([0.0, 0.0], v1_val)
+
+        # Run a step of AdagradDA
+        update.run()
+
+        v0_val, v1_val = self.evaluate([var0, var1])
+        # Let g be the gradient accumulator, gg be the gradient squared
+        # accumulator, T be the global step, lr be the learning rate,
+        # and k the initial gradient squared accumulator value.
+        # w = \dfrac{sign(-g)*lr*|g - l1*T|_{+}}{l2*T*lr + \sqrt{k+gg})}
+        # For -0.1*3.0*(0.1 - 0)/(0 + sqrt(0.1 + 0.1*0.1)) = -0.904534
+        # similarly for others.
+        self.assertAllCloseAccordingToType(
+            np.array([-0.904534, -1.603567]), v0_val)
+        self.assertAllCloseAccordingToType(
+            np.array([-0.094821, -0.189358]), v1_val)
+
+  def testAdagradDAWithoutRegularizationBasic1(self):
+    self.doTestAdagradDAwithoutRegularizationBasic1()
+
+  def testResourceAdagradDAWithoutRegularizationBasic1(self):
+    self.doTestAdagradDAwithoutRegularizationBasic1(use_resource=True)
+
+  @test_util.run_v1_only("loss needs to be callable in v2")
+  def testMinimizeSparseResourceVariable(self):
+    for dtype in [dtypes.float32, dtypes.float64]:
+      with self.cached_session():
+        var0 = resource_variable_ops.ResourceVariable([[1.0, 2.0]], dtype=dtype)
+        global_step = resource_variable_ops.ResourceVariable(
+            0, dtype=dtypes.int64)
+        x = constant_op.constant([[4.0], [5.0]], dtype=dtype)
+        pred = math_ops.matmul(embedding_ops.embedding_lookup([var0], [0]), x)
+        loss = pred * pred
+        sgd_op = adagrad_da.AdagradDAOptimizer(
+            1.0, global_step).minimize(loss)
+        self.evaluate(variables.global_variables_initializer())
+        # Fetch params to validate initial values
+        self.assertAllCloseAccordingToType([[1.0, 2.0]], self.evaluate(var0))
+        # Run 1 step of sgd
+        sgd_op.run()
+        # Validate updated params
+        self.assertAllCloseAccordingToType([[-1, -1]],
+                                           self.evaluate(var0),
+                                           rtol=0.01)
+
+  def testAdagradDAwithoutRegularizationBasic2(self):
+    for dtype in [dtypes.float64, dtypes.float32]:
+      with ops.Graph().as_default(), self.cached_session():
+        global_step = variables.Variable(0, dtype=dtypes.int64)
+        var0 = variables.Variable([1.0, 2.0], dtype=dtype)
+        var1 = variables.Variable([4.0, 3.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+
+        opt = adagrad_da.AdagradDAOptimizer(
+            3.0,
+            global_step,
+            initial_gradient_squared_accumulator_value=0.1,
+            l1_regularization_strength=0.0,
+            l2_regularization_strength=0.0)
+        update = opt.apply_gradients(
+            zip([grads0, grads1], [var0, var1]), global_step=global_step)
+        self.evaluate(variables.global_variables_initializer())
+
+        v0_val, v1_val = self.evaluate([var0, var1])
+        self.assertAllCloseAccordingToType([1.0, 2.0], v0_val)
+        self.assertAllCloseAccordingToType([4.0, 3.0], v1_val)
+
+        # Run a step of AdagradDA
+        update.run()
+
+        v0_val, v1_val = self.evaluate([var0, var1])
+        self.assertAllCloseAccordingToType(
+            np.array([-0.904534, -1.603567]), v0_val)
+        self.assertAllCloseAccordingToType(
+            np.array([-0.094821, -0.189358]), v1_val)
+
+  def testAdagradDAWithL1(self):
+    for dtype in [dtypes.float64, dtypes.float32]:
+      with ops.Graph().as_default(), self.cached_session():
+        global_step = variables.Variable(0, dtype=dtypes.int64)
+        var0 = variables.Variable([1.0, 2.0], dtype=dtype)
+        var1 = variables.Variable([4.0, 3.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+
+        opt = adagrad_da.AdagradDAOptimizer(
+            3.0,
+            global_step,
+            initial_gradient_squared_accumulator_value=0.1,
+            l1_regularization_strength=0.001,
+            l2_regularization_strength=0.0)
+        update = opt.apply_gradients(
+            zip([grads0, grads1], [var0, var1]), global_step=global_step)
+        self.evaluate(variables.global_variables_initializer())
+
+        v0_val, v1_val = self.evaluate([var0, var1])
+        self.assertAllCloseAccordingToType([1.0, 2.0], v0_val)
+        self.assertAllCloseAccordingToType([4.0, 3.0], v1_val)
+
+        # Run a step of AdagradDA
+        update.run()
+
+        v0_val, v1_val = self.evaluate([var0, var1])
+        self.assertAllCloseAccordingToType(
+            np.array([-0.895489, -1.59555]), v0_val)
+        self.assertAllCloseAccordingToType(
+            np.array([-0.085339, -0.17989]), v1_val)
+
+  def testAdagradDAWithL1_L2(self):
+    for dtype in [dtypes.float64, dtypes.float32]:
+      with ops.Graph().as_default(), self.cached_session():
+        global_step = variables.Variable(0, dtype=dtypes.int64)
+        var0 = variables.Variable([1.0, 2.0], dtype=dtype)
+        var1 = variables.Variable([4.0, 3.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+
+        opt = adagrad_da.AdagradDAOptimizer(
+            3.0,
+            global_step,
+            initial_gradient_squared_accumulator_value=0.1,
+            l1_regularization_strength=0.001,
+            l2_regularization_strength=2.0)
+        update = opt.apply_gradients(
+            zip([grads0, grads1], [var0, var1]), global_step=global_step)
+        self.evaluate(variables.global_variables_initializer())
+
+        v0_val, v1_val = self.evaluate([var0, var1])
+        self.assertAllCloseAccordingToType([1.0, 2.0], v0_val)
+        self.assertAllCloseAccordingToType([4.0, 3.0], v1_val)
+
+        # Run a step of AdagradDA
+        update.run()
+
+        v0_val, v1_val = self.evaluate([var0, var1])
+        self.assertAllCloseAccordingToType(
+            np.array([-0.046907, -0.093659]), v0_val)
+        self.assertAllCloseAccordingToType(
+            np.array([-0.004275, -0.009023]), v1_val)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/adagrad_test.py

@@ -0,0 +1,342 @@
+# 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 tests for aggregate operations."""
+
+import numpy as np
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import embedding_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import adagrad
+
+
+class AdagradOptimizerTest(test.TestCase):
+
+  def doTestBasic(self,
+                  use_locking=False,
+                  use_resource=False,
+                  use_callable_params=False):
+    for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+      if use_resource:
+        var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
+        var1 = resource_variable_ops.ResourceVariable([3.0, 4.0], dtype=dtype)
+      else:
+        var0 = variables.Variable([1.0, 2.0], dtype=dtype)
+        var1 = variables.Variable([3.0, 4.0], dtype=dtype)
+      grads0 = constant_op.constant([0.1, 0.1], dtype=dtype)
+      grads1 = constant_op.constant([0.01, 0.01], dtype=dtype)
+
+      learning_rate = lambda: 3.0
+      if not use_callable_params:
+        learning_rate = learning_rate()
+
+      ada_opt = adagrad.AdagradOptimizer(
+          learning_rate, initial_accumulator_value=0.1, use_locking=use_locking)
+
+      if not context.executing_eagerly():
+        ada_update = ada_opt.apply_gradients(
+            zip([grads0, grads1], [var0, var1]))
+        self.evaluate(variables.global_variables_initializer())
+
+      # Fetch params to validate initial values
+      v0_val, v1_val = self.evaluate([var0, var1])
+      self.assertAllClose([1.0, 2.0], v0_val)
+      self.assertAllClose([3.0, 4.0], v1_val)
+
+      # Run 3 steps of adagrad
+      for _ in range(3):
+        if not context.executing_eagerly():
+          self.evaluate(ada_update)
+        else:
+          ada_opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+
+      # Validate updated params
+      v0_val, v1_val = self.evaluate([var0, var1])
+      self.assertAllCloseAccordingToType(
+          np.array([-1.6026098728179932, -0.6026098728179932]), v0_val)
+      self.assertAllCloseAccordingToType(
+          np.array([2.715679168701172, 3.715679168701172]), v1_val)
+
+  def testBasic(self):
+    self.doTestBasic(use_locking=False)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testBasicResource(self):
+    self.doTestBasic(use_locking=False, use_resource=True)
+
+  def testBasicCallableParams(self):
+    with context.eager_mode():
+      self.doTestBasic(
+          use_locking=False, use_resource=True, use_callable_params=True)
+
+  def testBasicLocked(self):
+    self.doTestBasic(use_locking=True)
+
+  def testMinimizeSparseResourceVariable(self):
+    with ops.Graph().as_default():
+      for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+        with self.cached_session():
+          var0 = resource_variable_ops.ResourceVariable(
+              [[1.0, 2.0], [3.0, 4.0]], dtype=dtype)
+          x = constant_op.constant([[4.0], [5.0]], dtype=dtype)
+          pred = math_ops.matmul(embedding_ops.embedding_lookup([var0], [0]), x)
+          loss = pred * pred
+          sgd_op = adagrad.AdagradOptimizer(1.0).minimize(loss)
+          self.evaluate(variables.global_variables_initializer())
+          # Fetch params to validate initial values
+          self.assertAllCloseAccordingToType([[1.0, 2.0], [3.0, 4.0]],
+                                             self.evaluate(var0))
+          # Run 1 step of sgd
+          sgd_op.run()
+          # Validate updated params
+          self.assertAllCloseAccordingToType([[0, 1], [3, 4]],
+                                             self.evaluate(var0),
+                                             atol=0.01)
+
+  def testTensorLearningRate(self):
+    with ops.Graph().as_default():
+      for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+        with self.cached_session():
+          var0 = variables.Variable([1.0, 2.0], dtype=dtype)
+          var1 = variables.Variable([3.0, 4.0], dtype=dtype)
+          grads0 = constant_op.constant([0.1, 0.1], dtype=dtype)
+          grads1 = constant_op.constant([0.01, 0.01], dtype=dtype)
+          ada_opt = adagrad.AdagradOptimizer(
+              constant_op.constant(3.0), initial_accumulator_value=0.1)
+          ada_update = ada_opt.apply_gradients(
+              zip([grads0, grads1], [var0, var1]))
+          self.evaluate(variables.global_variables_initializer())
+          # Fetch params to validate initial values
+          self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+          self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+          # Run 3 steps of adagrad
+          for _ in range(3):
+            ada_update.run()
+          # Validate updated params
+          self.assertAllCloseAccordingToType(
+              np.array([-1.6026098728179932, -0.6026098728179932]),
+              self.evaluate(var0))
+          self.assertAllCloseAccordingToType(
+              np.array([2.715679168701172, 3.715679168701172]),
+              self.evaluate(var1))
+
+  def testSparseBasic(self):
+    with ops.Graph().as_default():
+      for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+        with self.cached_session():
+          var0 = variables.Variable([[1.0], [2.0]], dtype=dtype)
+          var1 = variables.Variable([[3.0], [4.0]], dtype=dtype)
+          grads0 = indexed_slices.IndexedSlices(
+              constant_op.constant(
+                  [0.1], shape=[1, 1], dtype=dtype),
+              constant_op.constant([0]),
+              constant_op.constant([2, 1]))
+          grads1 = indexed_slices.IndexedSlices(
+              constant_op.constant(
+                  [0.01], shape=[1, 1], dtype=dtype),
+              constant_op.constant([1]),
+              constant_op.constant([2, 1]))
+          ada_opt = adagrad.AdagradOptimizer(3.0, initial_accumulator_value=0.1)
+          ada_update = ada_opt.apply_gradients(
+              zip([grads0, grads1], [var0, var1]))
+          self.evaluate(variables.global_variables_initializer())
+          # Fetch params to validate initial values
+          self.assertAllClose([[1.0], [2.0]], self.evaluate(var0))
+          self.assertAllClose([[3.0], [4.0]], self.evaluate(var1))
+          # Run 3 step of sgd
+          for _ in range(3):
+            ada_update.run()
+          # Validate updated params
+          self.assertAllCloseAccordingToType(
+              np.array([[-1.6026098728179932], [2.0]]), self.evaluate(var0))
+          self.assertAllCloseAccordingToType(
+              np.array([[3.0], [3.715679168701172]]), self.evaluate(var1))
+
+  def testSparseRepeatedIndices(self):
+    with ops.Graph().as_default():
+      for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+        with self.cached_session():
+          repeated_index_update_var = variables.Variable(
+              [[1.0], [2.0]], dtype=dtype)
+          aggregated_update_var = variables.Variable(
+              [[1.0], [2.0]], dtype=dtype)
+          grad_repeated_index = indexed_slices.IndexedSlices(
+              constant_op.constant(
+                  [0.1, 0.1], shape=[2, 1], dtype=dtype),
+              constant_op.constant([1, 1]),
+              constant_op.constant([2, 1]))
+          grad_aggregated = indexed_slices.IndexedSlices(
+              constant_op.constant(
+                  [0.2], shape=[1, 1], dtype=dtype),
+              constant_op.constant([1]),
+              constant_op.constant([2, 1]))
+          repeated_update = adagrad.AdagradOptimizer(3.0).apply_gradients(
+              [(grad_repeated_index, repeated_index_update_var)])
+          aggregated_update = adagrad.AdagradOptimizer(3.0).apply_gradients(
+              [(grad_aggregated, aggregated_update_var)])
+          self.evaluate(variables.global_variables_initializer())
+          self.assertAllClose(aggregated_update_var,
+                              self.evaluate(repeated_index_update_var))
+          for _ in range(3):
+            repeated_update.run()
+            aggregated_update.run()
+            self.assertAllClose(aggregated_update_var,
+                                self.evaluate(repeated_index_update_var))
+
+  def testSparseRepeatedIndicesResourceVariable(self):
+    with ops.Graph().as_default():
+      for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+        with self.cached_session():
+          var_repeated = resource_variable_ops.ResourceVariable(
+              [1.0, 2.0], dtype=dtype)
+          loss_repeated = math_ops.reduce_sum(
+              embedding_ops.embedding_lookup(var_repeated, [0, 0]))
+          var_aggregated = resource_variable_ops.ResourceVariable(
+              [1.0, 2.0], dtype=dtype)
+          loss_aggregated = 2 * math_ops.reduce_sum(
+              embedding_ops.embedding_lookup(var_aggregated, [0]))
+          update_op_repeated = adagrad.AdagradOptimizer(
+              2.0).minimize(loss_repeated)
+          update_op_aggregated = adagrad.AdagradOptimizer(
+              2.0).minimize(loss_aggregated)
+          self.evaluate(variables.global_variables_initializer())
+          self.assertAllCloseAccordingToType(
+              self.evaluate(var_repeated), self.evaluate(var_aggregated))
+          for _ in range(3):
+            update_op_repeated.run()
+            update_op_aggregated.run()
+            self.assertAllCloseAccordingToType(
+                self.evaluate(var_repeated), self.evaluate(var_aggregated))
+
+  def testSparseStability(self):
+    with ops.Graph().as_default():
+      for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+        with self.cached_session():
+          shape = [1, 6]
+          var0 = variables.Variable(
+              [[
+                  0.00872496, -0.106952, 0.110467, 0.226505, -0.0147257,
+                  -0.0105945
+              ]],
+              dtype=dtype)
+          grads0 = indexed_slices.IndexedSlices(
+              constant_op.constant(
+                  [[
+                      -5.91278e-05, 5.31673e-05, -2.5779e-06, 4.29153e-05,
+                      -8.4877e-05, -9.48906e-05
+                  ]],
+                  shape=shape,
+                  dtype=dtype),
+              constant_op.constant([0]),
+              constant_op.constant(shape))
+          ada_opt = adagrad.AdagradOptimizer(1.0, initial_accumulator_value=0.1)
+          ada_update = ada_opt.apply_gradients(zip([grads0], [var0]))
+          self.assertEqual(["accumulator"], ada_opt.get_slot_names())
+          slot0 = ada_opt.get_slot(var0, "accumulator")
+          init = variables.global_variables_initializer()
+          for _ in range(100):
+            init.run()
+            ada_update.run()
+            self.assertAllCloseAccordingToType(
+                np.array([[0.1, 0.1, 0.1, 0.1, 0.1, 0.1]]),
+                self.evaluate(slot0))
+            self.assertAllCloseAccordingToType(
+                np.array([[
+                    0.00891194, -0.10712013, 0.11047515, 0.22636929, -0.0144573,
+                    -0.01029443
+                ]]), self.evaluate(var0))
+
+  def testSharing(self):
+    with ops.Graph().as_default():
+      for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+        with self.cached_session():
+          var0 = variables.Variable([1.0, 2.0], dtype=dtype)
+          var1 = variables.Variable([3.0, 4.0], dtype=dtype)
+          grads0 = constant_op.constant([0.1, 0.1], dtype=dtype)
+          grads1 = constant_op.constant([0.01, 0.01], dtype=dtype)
+          ada_opt = adagrad.AdagradOptimizer(3.0)
+          # Apply the optimizer twice.  Both applications will use
+          # the same accums.
+          ada_update1 = ada_opt.apply_gradients(
+              zip([grads0, grads1], [var0, var1]))
+          ada_update2 = ada_opt.apply_gradients(
+              zip([grads0, grads1], [var0, var1]))
+          self.assertEqual(["accumulator"], ada_opt.get_slot_names())
+          slot0 = ada_opt.get_slot(var0, "accumulator")
+          self.assertEqual(slot0.get_shape(), var0.get_shape())
+          slot1 = ada_opt.get_slot(var1, "accumulator")
+          self.assertEqual(slot1.get_shape(), var1.get_shape())
+          self.evaluate(variables.global_variables_initializer())
+
+          # Fetch params to validate initial values.
+          self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+          self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+          # Mix the first and the second adagrad for 3 steps.
+          ada_update1.run()
+          ada_update2.run()
+          ada_update1.run()
+          # Validate updated params (the same as with only 1 Adagrad).
+          self.assertAllCloseAccordingToType(
+              np.array([-1.6026098728179932, -0.6026098728179932]),
+              self.evaluate(var0))
+          self.assertAllCloseAccordingToType(
+              np.array([2.715679168701172, 3.715679168701172]),
+              self.evaluate(var1))
+
+  def testDynamicShapeVariableWithCallableInit(self):
+    with ops.Graph().as_default():
+      var0 = variable_scope.get_variable("var0",
+                                         initializer=constant_op.constant(1.),
+                                         validate_shape=False)
+
+      grads0 = constant_op.constant(0.1, dtype=dtypes.float32)
+      learning_rate = lambda: 3.0
+
+      ada_opt = adagrad.AdagradOptimizer(
+          learning_rate, initial_accumulator_value=0.1, use_locking=True)
+
+      if not context.executing_eagerly():
+        ada_update = ada_opt.apply_gradients(
+            zip([grads0], [var0]))
+        self.evaluate(variables.global_variables_initializer())
+
+      # Fetch params to validate initial values
+      v0_val = self.evaluate([var0])
+      self.assertAllClose([1.0], v0_val)
+
+      # Run 3 steps of adagrad
+      for _ in range(3):
+        if not context.executing_eagerly():
+          self.evaluate(ada_update)
+        else:
+          ada_opt.apply_gradients(zip([grads0], [var0]))
+
+      # Validate updated params
+      v0_val = self.evaluate([var0])
+      self.assertAllCloseAccordingToType(
+          np.array([-1.6026098728179932]), v0_val)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/adam.py

@@ -0,0 +1,306 @@
+# 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.
+# ==============================================================================
+"""Adam for TensorFlow."""
+from tensorflow.python.eager import context
+from tensorflow.python.framework import 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 resource_variable_ops
+from tensorflow.python.ops import state_ops
+from tensorflow.python.training import optimizer
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(v1=["train.AdamOptimizer"])
+class AdamOptimizer(optimizer.Optimizer):
+  """Optimizer that implements the Adam algorithm.
+
+  References:
+    Adam - A Method for Stochastic Optimization:
+      [Kingma et al., 2015](https://arxiv.org/abs/1412.6980)
+      ([pdf](https://arxiv.org/pdf/1412.6980.pdf))
+
+  @compatibility(TF2)
+  tf.compat.v1.train.AdamOptimizer is compatible with eager mode and
+  `tf.function`.
+  When eager execution is enabled, `learning_rate`, `beta1`, `beta2`, 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.
+
+  To switch to native TF2 style, use [`tf.keras.optimizers.Adam`]
+  (https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adam)
+  instead. Please notice that due to the implementation differences,
+  `tf.keras.optimizers.Adam` and
+  `tf.compat.v1.train.AdamOptimizer` may have slight differences in
+  floating point numerics even though the formula used for the variable
+  updates still matches.
+
+  #### Structural Mapping to Native TF2
+
+  Before:
+
+  ```python
+  optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=0.001)
+  ```
+
+  After:
+
+  ```python
+  optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
+  ```
+
+  #### How to Map Arguments
+  |TF1 Arg Name          |TF2 Arg Name |Note                  |
+  |----------------------|-------------|----------------------|
+  |learning_rate         |learning_rate|Be careful of setting learning_rate as a
+  :                      :             : tensor value computed from the global
+  :                      :             : step. In TF1 this was usually meant to
+  :                      :             : imply a dynamic learning rate and would
+  :                      :             : recompute in each step. In TF2 (eager +
+  :                      :             : function) it will treat it as a scalar
+  :                      :             : value that only gets computed once
+  :                      :             : instead of a symbolic placeholder to be
+  :                      :             : computed each time.                   :
+  |beta1                 |beta_1        |                      |
+  |beta2                 |beta_2        |                      |
+  |epsilon               |epsilon      | Default value is 1e-08 in TF1, but
+  :                      :             : 1e-07 in TF2.                     :
+  |use_locking           |N/A          |Not applicable in TF2. |
+
+  #### Before & After Usage Example
+  Before:
+
+  ```python
+  x = tf.Variable([1,2,3], dtype=tf.float32)
+  grad = tf.constant([0.1, 0.2, 0.3])
+  optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=0.001)
+  optimizer.apply_gradients(zip([grad], [x]))
+  ```
+
+  After:
+
+  ```python
+  x = tf.Variable([1,2,3], dtype=tf.float32)
+  grad = tf.constant([0.1, 0.2, 0.3])
+  optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
+  optimizer.apply_gradients(zip([grad], [x]))
+  ```
+
+  @end_compatibility
+  """
+
+  def __init__(
+      self,
+      learning_rate=0.001,
+      beta1=0.9,
+      beta2=0.999,
+      epsilon=1e-8,
+      use_locking=False,
+      name="Adam",
+      use_own_namescope_for_non_slot_vars=False,
+  ):
+    r"""Construct a new Adam optimizer.
+
+    Initialization:
+
+    $$m_0 := 0 \text{(Initialize initial 1st moment vector)}$$
+    $$v_0 := 0 \text{(Initialize initial 2nd moment vector)}$$
+    $$t := 0 \text{(Initialize timestep)}$$
+
+    The update rule for `variable` with gradient `g` uses an optimization
+    described at the end of section 2 of the paper:
+
+    $$t := t + 1$$
+    $$\text{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 * g$$
+    $$\text{variable} := \text{variable} -
+      \text{lr}_t * m_t / (\sqrt{v_t} + \epsilon)$$
+
+    The default value of 1e-8 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 AdamOptimizer 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).
+
+    Args:
+      learning_rate: A Tensor or a floating point value.  The learning rate.
+      beta1: A float value or a constant float tensor. The exponential decay
+        rate for the 1st moment estimates.
+      beta2: A float value or a constant float tensor. The exponential decay
+        rate for the 2nd moment estimates.
+      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.
+      use_locking: If True use locks for update operations.
+      name: Optional name for the operations created when applying gradients.
+        Defaults to "Adam".
+      use_own_namescope_for_non_slot_vars: If True, use a root namescope for
+        non-slot variables.
+    """
+
+    super().__init__(use_locking, name, use_own_namescope_for_non_slot_vars)
+    self._lr = learning_rate
+    self._beta1 = beta1
+    self._beta2 = beta2
+    self._epsilon = epsilon
+
+    # Tensor versions of the constructor arguments, created in _prepare().
+    self._lr_t = None
+    self._beta1_t = None
+    self._beta2_t = None
+    self._epsilon_t = None
+
+  def _get_beta_accumulators(self):
+    with ops.init_scope():
+      if context.executing_eagerly():
+        graph = None
+      else:
+        graph = ops.get_default_graph()
+      return (self._get_non_slot_variable("beta1_power", graph=graph),
+              self._get_non_slot_variable("beta2_power", graph=graph))
+
+  def _create_slots(self, var_list):
+    # Create the beta1 and beta2 accumulators on the same device as the first
+    # variable. Sort the var_list to make sure this device is consistent across
+    # workers (these need to go on the same PS, otherwise some updates are
+    # silently ignored).
+    first_var = min(var_list, key=lambda x: x.name)
+    self._create_non_slot_variable(
+        initial_value=self._beta1, name="beta1_power", colocate_with=first_var)
+    self._create_non_slot_variable(
+        initial_value=self._beta2, name="beta2_power", colocate_with=first_var)
+
+    # Create slots for the first and second moments.
+    for v in var_list:
+      self._zeros_slot(v, "m", self._name)
+      self._zeros_slot(v, "v", self._name)
+
+  def _prepare(self):
+    lr = self._call_if_callable(self._lr)
+    beta1 = self._call_if_callable(self._beta1)
+    beta2 = self._call_if_callable(self._beta2)
+    epsilon = self._call_if_callable(self._epsilon)
+
+    self._lr_t = ops.convert_to_tensor(lr, name="learning_rate")
+    self._beta1_t = ops.convert_to_tensor(beta1, name="beta1")
+    self._beta2_t = ops.convert_to_tensor(beta2, name="beta2")
+    self._epsilon_t = ops.convert_to_tensor(epsilon, name="epsilon")
+
+  def _apply_dense(self, grad, var):
+    m = self.get_slot(var, "m")
+    v = self.get_slot(var, "v")
+    beta1_power, beta2_power = self._get_beta_accumulators()
+    return gen_training_ops.apply_adam(
+        var,
+        m,
+        v,
+        math_ops.cast(beta1_power, var.dtype.base_dtype),
+        math_ops.cast(beta2_power, var.dtype.base_dtype),
+        math_ops.cast(self._lr_t, var.dtype.base_dtype),
+        math_ops.cast(self._beta1_t, var.dtype.base_dtype),
+        math_ops.cast(self._beta2_t, var.dtype.base_dtype),
+        math_ops.cast(self._epsilon_t, var.dtype.base_dtype),
+        grad,
+        use_locking=self._use_locking).op
+
+  def _resource_apply_dense(self, grad, var):
+    m = self.get_slot(var, "m")
+    v = self.get_slot(var, "v")
+    beta1_power, beta2_power = self._get_beta_accumulators()
+    return gen_training_ops.resource_apply_adam(
+        var.handle,
+        m.handle,
+        v.handle,
+        math_ops.cast(beta1_power, grad.dtype.base_dtype),
+        math_ops.cast(beta2_power, grad.dtype.base_dtype),
+        math_ops.cast(self._lr_t, grad.dtype.base_dtype),
+        math_ops.cast(self._beta1_t, grad.dtype.base_dtype),
+        math_ops.cast(self._beta2_t, grad.dtype.base_dtype),
+        math_ops.cast(self._epsilon_t, grad.dtype.base_dtype),
+        grad,
+        use_locking=self._use_locking)
+
+  def _apply_sparse_shared(self, grad, var, indices, scatter_add):
+    beta1_power, beta2_power = self._get_beta_accumulators()
+    beta1_power = math_ops.cast(beta1_power, var.dtype.base_dtype)
+    beta2_power = math_ops.cast(beta2_power, var.dtype.base_dtype)
+    lr_t = math_ops.cast(self._lr_t, var.dtype.base_dtype)
+    beta1_t = math_ops.cast(self._beta1_t, var.dtype.base_dtype)
+    beta2_t = math_ops.cast(self._beta2_t, var.dtype.base_dtype)
+    epsilon_t = math_ops.cast(self._epsilon_t, var.dtype.base_dtype)
+    lr = (lr_t * math_ops.sqrt(1 - beta2_power) / (1 - beta1_power))
+    # m_t = beta1 * m + (1 - beta1) * g_t
+    m = self.get_slot(var, "m")
+    m_scaled_g_values = grad * (1 - beta1_t)
+    m_t = state_ops.assign(m, m * beta1_t, use_locking=self._use_locking)
+    with ops.control_dependencies([m_t]):
+      m_t = 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) * (1 - beta2_t)
+    v_t = state_ops.assign(v, v * beta2_t, use_locking=self._use_locking)
+    with ops.control_dependencies([v_t]):
+      v_t = scatter_add(v, indices, v_scaled_g_values)
+    v_sqrt = math_ops.sqrt(v_t)
+    var_update = state_ops.assign_sub(
+        var, lr * m_t / (v_sqrt + epsilon_t), use_locking=self._use_locking)
+    return control_flow_ops.group(*[var_update, m_t, v_t])
+
+  def _apply_sparse(self, grad, var):
+    return self._apply_sparse_shared(
+        grad.values,
+        var,
+        grad.indices,
+        lambda x, i, v: state_ops.scatter_add(  # pylint: disable=g-long-lambda
+            x,
+            i,
+            v,
+            use_locking=self._use_locking))
+
+  def _resource_scatter_add(self, x, i, v):
+    with ops.control_dependencies(
+        [resource_variable_ops.resource_scatter_add(x.handle, i, v)]):
+      return x.value()
+
+  def _resource_apply_sparse(self, grad, var, indices):
+    return self._apply_sparse_shared(grad, var, indices,
+                                     self._resource_scatter_add)
+
+  def _finish(self, update_ops, name_scope):
+    # Update the power accumulators.
+    with ops.control_dependencies(update_ops):
+      beta1_power, beta2_power = self._get_beta_accumulators()
+      with ops.colocate_with(beta1_power):
+        update_beta1 = beta1_power.assign(
+            beta1_power * self._beta1_t, use_locking=self._use_locking)
+        update_beta2 = beta2_power.assign(
+            beta2_power * self._beta2_t, use_locking=self._use_locking)
+    return control_flow_ops.group(
+        *update_ops + [update_beta1, update_beta2], name=name_scope)

tensorflow/python/training/adam_test.py

@@ -0,0 +1,434 @@
+# 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.
+# ==============================================================================
+"""Tests for Adam."""
+
+import numpy as np
+
+from tensorflow.python.client import session
+from tensorflow.python.compiler.xla.experimental import xla_sharding
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import ref_variable
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import adam
+
+
+def adam_update_numpy(param,
+                      g_t,
+                      t,
+                      m,
+                      v,
+                      alpha=0.001,
+                      beta1=0.9,
+                      beta2=0.999,
+                      epsilon=1e-8):
+  alpha_t = alpha * np.sqrt(1 - beta2**t) / (1 - beta1**t)
+
+  m_t = beta1 * m + (1 - beta1) * g_t
+  v_t = beta2 * v + (1 - beta2) * g_t * g_t
+
+  param_t = param - alpha_t * m_t / (np.sqrt(v_t) + epsilon)
+  return param_t, m_t, v_t
+
+
+class AdamOptimizerTest(test.TestCase):
+
+  def doTestSparse(self, use_resource=False):
+    for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+      with self.cached_session():
+        # Initialize variables for numpy implementation.
+        m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0
+        var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
+        grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype)
+        var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
+        grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype)
+
+        if use_resource:
+          var0 = resource_variable_ops.ResourceVariable(var0_np)
+          var1 = resource_variable_ops.ResourceVariable(var1_np)
+        else:
+          var0 = ref_variable.RefVariable(var0_np)
+          var1 = ref_variable.RefVariable(var1_np)
+        grads0_np_indices = np.array([0, 1], dtype=np.int32)
+        grads0 = indexed_slices.IndexedSlices(
+            constant_op.constant(grads0_np),
+            constant_op.constant(grads0_np_indices), constant_op.constant([2]))
+        grads1_np_indices = np.array([0, 1], dtype=np.int32)
+        grads1 = indexed_slices.IndexedSlices(
+            constant_op.constant(grads1_np),
+            constant_op.constant(grads1_np_indices), constant_op.constant([2]))
+        opt = adam.AdamOptimizer()
+        update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+        self.evaluate(variables.global_variables_initializer())
+
+        # Fetch params to validate initial values
+        self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+        self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+
+        beta1_power, beta2_power = opt._get_beta_accumulators()
+
+        # Run 3 steps of Adam
+        for t in range(1, 4):
+          self.assertAllCloseAccordingToType(0.9**t, self.evaluate(beta1_power))
+          self.assertAllCloseAccordingToType(0.999**t,
+                                             self.evaluate(beta2_power))
+          update.run()
+
+          var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0)
+          var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1)
+
+          # Validate updated params
+          self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0))
+          self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1))
+
+  def testSparse(self):
+    with ops.Graph().as_default():
+      self.doTestSparse(use_resource=False)
+
+  def testResourceSparse(self):
+    with ops.Graph().as_default():
+      self.doTestSparse(use_resource=True)
+
+  def testSparseDevicePlacement(self):
+    with ops.Graph().as_default():
+      for index_dtype in [dtypes.int32, dtypes.int64]:
+        with self.cached_session(force_gpu=test.is_gpu_available()):
+          # If a GPU is available, tests that all optimizer ops can be placed on
+          # it (i.e. they have GPU kernels).
+          var = variables.Variable([[1.0], [2.0]])
+          indices = constant_op.constant([0, 1], dtype=index_dtype)
+          gathered_sum = math_ops.reduce_sum(array_ops.gather(var, indices))
+          optimizer = adam.AdamOptimizer(3.0)
+          minimize_op = optimizer.minimize(gathered_sum)
+          self.evaluate(variables.global_variables_initializer())
+          minimize_op.run()
+
+  def testGatherGradientWithBadIndicesPolicy(self):
+    with ops.Graph().as_default():
+      with self.cached_session(force_gpu=test.is_gpu_available()):
+        var = variables.Variable([1.0, 2.0])
+        indices = constant_op.constant([[1], [-1], [0]], dtype=dtypes.int32)
+        out = array_ops.gather_nd(var,
+                                  array_ops.expand_dims(indices, axis=-1),
+                                  batch_dims=0,
+                                  bad_indices_policy="IGNORE")
+        optimizer = adam.AdamOptimizer(2.0, 0.0, 1.0)
+        minimize_op = optimizer.minimize(out)
+        self.evaluate(variables.global_variables_initializer())
+        minimize_op.run()
+
+  def testSparseRepeatedIndices(self):
+    with ops.Graph().as_default():
+      for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+        with self.cached_session():
+          repeated_index_update_var = variables.Variable(
+              [[1.0], [2.0]], dtype=dtype)
+          aggregated_update_var = variables.Variable(
+              [[1.0], [2.0]], dtype=dtype)
+          grad_repeated_index = indexed_slices.IndexedSlices(
+              constant_op.constant(
+                  [0.1, 0.1], shape=[2, 1], dtype=dtype),
+              constant_op.constant([1, 1]),
+              constant_op.constant([2, 1]))
+          grad_aggregated = indexed_slices.IndexedSlices(
+              constant_op.constant(
+                  [0.2], shape=[1, 1], dtype=dtype),
+              constant_op.constant([1]),
+              constant_op.constant([2, 1]))
+          repeated_update = adam.AdamOptimizer().apply_gradients(
+              [(grad_repeated_index, repeated_index_update_var)])
+          aggregated_update = adam.AdamOptimizer().apply_gradients(
+              [(grad_aggregated, aggregated_update_var)])
+          self.evaluate(variables.global_variables_initializer())
+          self.assertAllClose(aggregated_update_var,
+                              self.evaluate(repeated_index_update_var))
+          for _ in range(3):
+            repeated_update.run()
+            aggregated_update.run()
+            self.assertAllClose(aggregated_update_var,
+                                self.evaluate(repeated_index_update_var))
+
+  def doTestBasic(self, use_resource=False, use_callable_params=False):
+    if context.executing_eagerly() and not use_resource:
+      self.skipTest(
+          "Skipping test with use_resource=False and executing eagerly.")
+    for i, dtype in enumerate([dtypes.half, dtypes.float32, dtypes.float64]):
+      with self.session(graph=ops.Graph()):
+        # Initialize variables for numpy implementation.
+        m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0
+        var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
+        grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype)
+        var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
+        grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype)
+
+        if use_resource:
+          var0 = resource_variable_ops.ResourceVariable(
+              var0_np, name="var0_%d" % i)
+          var1 = resource_variable_ops.ResourceVariable(
+              var1_np, name="var1_%d" % i)
+        else:
+          var0 = ref_variable.RefVariable(var0_np)
+          var1 = ref_variable.RefVariable(var1_np)
+        grads0 = constant_op.constant(grads0_np)
+        grads1 = constant_op.constant(grads1_np)
+
+        learning_rate = lambda: 0.001
+        beta1 = lambda: 0.9
+        beta2 = lambda: 0.999
+        epsilon = lambda: 1e-8
+        if not use_callable_params:
+          learning_rate = learning_rate()
+          beta1 = beta1()
+          beta2 = beta2()
+          epsilon = epsilon()
+
+        opt = adam.AdamOptimizer(learning_rate=learning_rate)
+        update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+        opt_variables = opt.variables()
+        beta1_power, beta2_power = opt._get_beta_accumulators()
+        self.assertTrue(beta1_power is not None)
+        self.assertTrue(beta2_power is not None)
+        self.assertIn(beta1_power, opt_variables)
+        self.assertIn(beta2_power, opt_variables)
+        # Ensure that non-slot variables are the same type as the requested
+        # variables.
+        self.assertEqual(
+            use_resource,
+            resource_variable_ops.is_resource_variable(beta1_power))
+        self.assertEqual(
+            use_resource,
+            resource_variable_ops.is_resource_variable(beta2_power))
+
+        if not context.executing_eagerly():
+          with ops.Graph().as_default():
+            # Shouldn't return non-slot variables from other graphs.
+            self.assertEqual(0, len(opt.variables()))
+          self.evaluate(variables.global_variables_initializer())
+          # Fetch params to validate initial values
+          self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+          self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+
+        beta1_power, beta2_power = opt._get_beta_accumulators()
+
+        # Run 3 steps of Adam
+        for t in range(1, 4):
+          if not context.executing_eagerly():
+            self.evaluate(update)
+          elif t > 1:
+            opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+
+          self.assertAllCloseAccordingToType(0.9**(t + 1),
+                                             self.evaluate(beta1_power))
+          self.assertAllCloseAccordingToType(0.999**(t + 1),
+                                             self.evaluate(beta2_power))
+
+          var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0)
+          var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1)
+
+          # Validate updated params
+          self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0))
+          self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1))
+          if use_resource:
+            self.assertEqual("var0_%d/Adam:0" % (i,),
+                             opt.get_slot(var=var0, name="m").name)
+
+  def testBasic(self):
+    with self.cached_session():
+      self.doTestBasic(use_resource=False)
+
+  @test_util.run_in_graph_and_eager_modes
+  @test_util.disable_tfrt("b/168527439: invalid runtime fallback "
+                          "resource variable reference on GPU.")
+  def testResourceBasic(self):
+    self.doTestBasic(use_resource=True)
+
+  @test_util.disable_tfrt("b/153089059: cannot create half tensor on GPU.")
+  def testBasicCallableParams(self):
+    with context.eager_mode():
+      self.doTestBasic(use_resource=True, use_callable_params=True)
+
+  def testTensorLearningRate(self):
+    with ops.Graph().as_default():
+      for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+        with self.cached_session():
+          # Initialize variables for numpy implementation.
+          m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0
+          var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
+          grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype)
+          var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
+          grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype)
+
+          var0 = variables.Variable(var0_np)
+          var1 = variables.Variable(var1_np)
+          grads0 = constant_op.constant(grads0_np)
+          grads1 = constant_op.constant(grads1_np)
+          opt = adam.AdamOptimizer(constant_op.constant(0.001))
+          update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+          self.evaluate(variables.global_variables_initializer())
+
+          # Fetch params to validate initial values
+          self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+          self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+
+          beta1_power, beta2_power = opt._get_beta_accumulators()
+
+          # Run 3 steps of Adam
+          for t in range(1, 4):
+            self.assertAllCloseAccordingToType(0.9**t,
+                                               self.evaluate(beta1_power))
+            self.assertAllCloseAccordingToType(0.999**t,
+                                               self.evaluate(beta2_power))
+            update.run()
+
+            var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0)
+            var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1)
+
+            # Validate updated params
+            self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0))
+            self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1))
+
+  def testSharing(self):
+    with ops.Graph().as_default():
+      for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+        with self.cached_session():
+          # Initialize variables for numpy implementation.
+          m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0
+          var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
+          grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype)
+          var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
+          grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype)
+
+          var0 = variables.Variable(var0_np)
+          var1 = variables.Variable(var1_np)
+          grads0 = constant_op.constant(grads0_np)
+          grads1 = constant_op.constant(grads1_np)
+          opt = adam.AdamOptimizer()
+          update1 = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+          update2 = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+          self.evaluate(variables.global_variables_initializer())
+
+          beta1_power, beta2_power = opt._get_beta_accumulators()
+
+          # Fetch params to validate initial values
+          self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+          self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+
+          # Run 3 steps of intertwined Adam1 and Adam2.
+          for t in range(1, 4):
+            self.assertAllCloseAccordingToType(0.9**t,
+                                               self.evaluate(beta1_power))
+            self.assertAllCloseAccordingToType(0.999**t,
+                                               self.evaluate(beta2_power))
+            if t % 2 == 0:
+              update1.run()
+            else:
+              update2.run()
+
+            var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0)
+            var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1)
+
+            # Validate updated params
+            self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0))
+            self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1))
+
+  @test_util.disable_tfrt("b/168527439: invalid runtime fallback "
+                          "resource variable reference on GPU.")
+  def testTwoSessions(self):
+    optimizer = adam.AdamOptimizer()
+
+    with context.eager_mode():
+      var0 = variables.Variable(np.array([1.0, 2.0]), name="v0")
+      grads0 = constant_op.constant(np.array([0.1, 0.1]))
+      optimizer.apply_gradients([(grads0, var0)])
+
+    g = ops.Graph()
+    with g.as_default():
+      with session.Session():
+        var0 = variables.Variable(np.array([1.0, 2.0]), name="v0")
+        grads0 = constant_op.constant(np.array([0.1, 0.1]))
+        optimizer.apply_gradients([(grads0, var0)])
+
+    gg = ops.Graph()
+    with gg.as_default():
+      with session.Session():
+        var0 = variables.Variable(np.array([1.0, 2.0]), name="v0")
+        grads0 = constant_op.constant(np.array([0.1, 0.1]))
+
+        # If the optimizer saves any state not keyed by graph the following line
+        # fails.
+        optimizer.apply_gradients([(grads0, var0)])
+
+  @test_util.disable_tfrt("b/168527439: invalid runtime fallback "
+                          "resource variable reference on GPU.")
+  def testSlotsUniqueEager(self):
+    with context.eager_mode():
+      v1 = resource_variable_ops.ResourceVariable(1.)
+      v2 = resource_variable_ops.ResourceVariable(1.)
+      opt = adam.AdamOptimizer(1.)
+      opt.minimize(lambda: v1 + v2)
+      # There should be two non-slot variables, and two unique slot variables
+      # for v1 and v2 respectively.
+      self.assertEqual(6, len({id(v) for v in opt.variables()}))
+
+  @test_util.deprecated_graph_mode_only
+  def testXlaSharding(self):
+    dtype = dtypes.float32
+    with self.session(graph=ops.Graph()):
+      # Initialize variables for numpy implementation.
+      var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
+      grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype)
+      var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
+      grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype)
+
+      var0 = resource_variable_ops.ResourceVariable(var0_np, name="var0")
+      var1 = resource_variable_ops.ResourceVariable(var1_np, name="var1")
+      var0, var1 = [
+          xla_sharding.mesh_split(
+              v, np.array([0, 1]), [0], use_sharding_op=False)
+          for v in (var0, var1)
+      ]
+      grads0 = constant_op.constant(grads0_np)
+      grads1 = constant_op.constant(grads1_np)
+
+      learning_rate = lambda: 0.001
+
+      opt = adam.AdamOptimizer(learning_rate=learning_rate)
+      update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+
+      self.evaluate(variables.global_variables_initializer())
+      self.evaluate(update)
+      # The beta accumulators are not sharded.
+      beta1_power, beta2_power = opt._get_beta_accumulators()
+      self.assertIsNone(xla_sharding.get_tensor_sharding(beta1_power))
+      self.assertIsNone(xla_sharding.get_tensor_sharding(beta2_power))
+
+      # Variables and slots are sharded.
+      for v in (var0, var1):
+        self.assertIsNotNone(xla_sharding.get_tensor_sharding(v))
+        for slot_name in ("m", "v"):
+          slot = opt.get_slot(v, slot_name)
+          self.assertIsNotNone(xla_sharding.get_tensor_sharding(slot))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/basic_loops.py

@@ -0,0 +1,61 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Basic loop for training."""
+from tensorflow.python.framework import errors
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(v1=["train.basic_train_loop"])
+def basic_train_loop(supervisor,
+                     train_step_fn,
+                     args=None,
+                     kwargs=None,
+                     master=""):
+  """Basic loop to train a model.
+
+  Calls `train_step_fn` in a loop to train a model.  The function is called as:
+
+  ```python
+  train_step_fn(session, *args, **kwargs)
+  ```
+
+  It is passed a `tf.compat.v1.Session` in addition to `args` and `kwargs`.  The
+  function
+  typically runs one training step in the session.
+
+  Args:
+    supervisor: `tf.compat.v1.train.Supervisor` to run the training services.
+    train_step_fn: Callable to execute one training step.  Called repeatedly as
+      `train_step_fn(session, *args **kwargs)`.
+    args: Optional positional arguments passed to `train_step_fn`.
+    kwargs: Optional keyword arguments passed to `train_step_fn`.
+    master: Master to use to create the training session.  Defaults to `""`
+      which causes the session to be created in the local process.
+  """
+  if args is None:
+    args = []
+  if kwargs is None:
+    kwargs = {}
+  should_retry = True
+  while should_retry:
+    try:
+      should_retry = False
+      with supervisor.managed_session(master) as sess:
+        while not supervisor.should_stop():
+          train_step_fn(sess, *args, **kwargs)
+    except errors.AbortedError:
+      # Always re-run on AbortedError as it indicates a restart of one of the
+      # distributed tensorflow servers.
+      should_retry = True

tensorflow/python/training/basic_loops_test.py

@@ -0,0 +1,97 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Tests for basic_loops.py."""
+
+import os
+import shutil
+
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.framework import ops
+from tensorflow.python.platform import test
+from tensorflow.python.training import basic_loops
+from tensorflow.python.training import supervisor
+
+
+def _test_dir(test_name):
+  test_dir = os.path.join(test.get_temp_dir(), test_name)
+  if os.path.exists(test_dir):
+    shutil.rmtree(test_dir)
+  return test_dir
+
+
+class BasicTrainLoopTest(test.TestCase):
+
+  def testBasicTrainLoop(self):
+    logdir = _test_dir("basic_train_loop")
+    # Counts the number of calls.
+    num_calls = [0]
+
+    def train_fn(unused_sess, sv, y, a):
+      num_calls[0] += 1
+      self.assertEqual("y", y)
+      self.assertEqual("A", a)
+      if num_calls[0] == 3:
+        sv.request_stop()
+
+    with ops.Graph().as_default():
+      sv = supervisor.Supervisor(logdir=logdir)
+      basic_loops.basic_train_loop(
+          sv, train_fn, args=(sv, "y"), kwargs={"a": "A"})
+      self.assertEqual(3, num_calls[0])
+
+  def testBasicTrainLoopExceptionAborts(self):
+    logdir = _test_dir("basic_train_loop_exception_aborts")
+
+    def train_fn(unused_sess):
+      train_fn.counter += 1
+      if train_fn.counter == 3:
+        raise RuntimeError("Failed")
+
+    # Function attribute use to count the number of calls.
+    train_fn.counter = 0
+
+    with ops.Graph().as_default():
+      sv = supervisor.Supervisor(logdir=logdir)
+      with self.assertRaisesRegex(RuntimeError, "Failed"):
+        basic_loops.basic_train_loop(sv, train_fn)
+
+  def testBasicTrainLoopRetryOnAborted(self):
+    logdir = _test_dir("basic_train_loop_exception_aborts")
+
+    class AbortAndRetry:
+
+      def __init__(self):
+        self.num_calls = 0
+        self.retries_left = 2
+
+      def train_fn(self, unused_sess):
+        self.num_calls += 1
+        if self.num_calls % 3 == 2:
+          self.retries_left -= 1
+        if self.retries_left > 0:
+          raise errors_impl.AbortedError(None, None, "Aborted here")
+        else:
+          raise RuntimeError("Failed Again")
+
+    with ops.Graph().as_default():
+      sv = supervisor.Supervisor(logdir=logdir)
+      aar = AbortAndRetry()
+      with self.assertRaisesRegex(RuntimeError, "Failed Again"):
+        basic_loops.basic_train_loop(sv, aar.train_fn)
+      self.assertEqual(0, aar.retries_left)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/checkpoint_management.py

@@ -0,0 +1,26 @@
+# 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.
+# ==============================================================================
+
+# pylint: disable=invalid-name
+"""Save and restore variables."""
+
+
+# TODO(kathywu): Delete this file after all imports have been moved to the path
+# below.
+from tensorflow.python.checkpoint import checkpoint_management
+from tensorflow.python.util import deprecation
+
+__getattr__ = deprecation.deprecate_moved_module(
+    __name__, checkpoint_management, "2.9")

tensorflow/python/training/checkpoint_ops.py

@@ -0,0 +1,482 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Operations for generating and loading vocab remappings."""
+import math
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_checkpoint_ops
+from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import math_ops
+
+ops.NotDifferentiable("GenerateVocabRemapping")
+ops.NotDifferentiable("LoadAndRemapMatrix")
+
+
+def _load_and_remap_matrix(ckpt_path,
+                           old_tensor_name,
+                           new_row_vocab_offset,
+                           num_rows_to_load,
+                           new_col_vocab_size,
+                           initializer,
+                           old_row_vocab_size=-1,
+                           old_row_vocab_file=None,
+                           new_row_vocab_file=None,
+                           old_col_vocab_file=None,
+                           new_col_vocab_file=None,
+                           num_row_oov_buckets=0,
+                           num_col_oov_buckets=0,
+                           max_rows_in_memory=-1):
+  """Loads a 2-D (matrix) `Tensor` from checkpoint.
+
+  Generates 1D-remappings for rows and columns using the
+  `GenerateVocabRemapping` op, and initializes any anticipated values with the
+  provided initializer. Then, uses the `LoadAndRemapMatrix` op to create a
+  matrix that loads existing values from the checkpoint, while filling out
+  "missing" values with the newly initialized values. See
+  contrib/framework/ops/checkpoint_ops.cc for more information on the wrapped
+  functionality (LoadAndRemapMatrix). This wrapper can be used to perform only
+  row remapping or only col remapping. If only row remapping is desired,
+  {new,old}_col_vocab_file should be `None`, and vice versa for column
+  remapping.
+
+  NOTE: This only supports div-partitioning the vocabulary on the 1st dimension
+  (row axis) via `new_row_vocab_offset`.
+
+  Args:
+    ckpt_path: Path to the TensorFlow checkpoint (version 2, `TensorBundle`)
+      from which the old matrix `Tensor` will be loaded.
+    old_tensor_name: Name of the 2-D `Tensor` to load from checkpoint.
+    new_row_vocab_offset: A 0-indexed integer representing what line to
+      start reading at in the new row vocabulary. Used for partitioned
+      variables.
+    num_rows_to_load: Number of rows to load for the new vocabulary (note: to
+      support variable partitioning and partial loading, this does not need to
+      be the same as the number of entries in `new_row_vocab_file`).
+    new_col_vocab_size: Number of columns to load - should be the same as the
+      number of entries in `new_col_vocab_file`, since we don't support
+      partitioning along the column axis.
+    initializer: Callable initializer function that accepts a 1-D tensor as the
+      arg to specify the shape of the returned tensor. Used to initialize
+      missing values.
+    old_row_vocab_size: The number of entries to consider in the old vocabulary.
+      With the default value of -1, the entire old row vocabulary file will be
+      used.  Otherwise, only the first `old_row_vocab_size` entries will be
+      considered for remapping.Must be smaller than the length of
+      `old_row_vocab_file`.  NOTE: we do not provide an equivalent
+      `old_col_vocab_size` for classes.
+    old_row_vocab_file: A scalar `Tensor` of type `string` containing the
+      path to the old row vocabulary file. Can be None, which represents no
+      remapping on the row axis.
+    new_row_vocab_file: A scalar `Tensor` of type `string` containing the path
+      to the new row vocabulary file. Can be None, which represents no remapping
+      on the row axis - in which case, `new_row_vocab_offset` and
+      `num_rows_to_load` work under the assumption that the new row vocab is the
+      same as the old row vocab.
+    old_col_vocab_file: A scalar `Tensor` of type `string` containing the
+      path to the old column vocabulary file. Can be None, which represents no
+      remapping on the column axis.
+    new_col_vocab_file: A scalar `Tensor` of type `string` containing the path
+      to the new column vocabulary file. Can be None, which represents no
+      remapping on the column axis - in which case, `new_col_vocab_size` works
+      under the assumption that the new col vocab is the same as the old col
+      vocab.
+    num_row_oov_buckets: `int` specifying the number of out-of-vocabulary rows
+      to append. Must be >= 0.
+    num_col_oov_buckets: `int` specifying the number of out-of-vocabulary
+      columns to append. Must be >= 0.
+    max_rows_in_memory: `int` specifying the maximum number of rows to load from
+      the checkpoint at once. If less than or equal to 0, the entire matrix will
+      be loaded into memory. Setting this arg trades increased disk reads for
+      lower memory usage.
+
+  Returns:
+    A Tensor of shape `[num_rows_to_load + num_row_oov_buckets,
+    new_col_vocab_size + num_col_oov_buckets]`, with values loaded from the
+    specified tensor in the checkpoint, and any missing or OOV values
+    initialized with the given `initializer`.
+
+  Raises:
+    ValueError: If `num_row_oov_buckets` or `num_col_oov_buckets` < 0.
+    ValueError: If either `old_row_vocab_file` or `new_row_vocab_file` is
+      provided, while the other is not. Same for `old_col_vocab_file` and
+      `new_col_vocab_file`.
+    ValueError: If neither row vocabs or col vocabs are provided.
+  """
+  if num_row_oov_buckets < 0:
+    raise ValueError("num_row_oov_buckets must be >= 0, but received %d" %
+                     num_row_oov_buckets)
+  if num_col_oov_buckets < 0:
+    raise ValueError("num_col_oov_buckets must be >= 0, but received %d" %
+                     num_col_oov_buckets)
+
+  if bool(old_row_vocab_file) != bool(new_row_vocab_file):
+    raise ValueError(
+        "old_row_vocab_file and new_row_vocab_file must both be specified or "
+        "left unspecified. old_row_vocab_file='{}', new_row_vocab_file='{}'".
+        format(old_row_vocab_file, new_row_vocab_file))
+  if bool(old_col_vocab_file) != bool(new_col_vocab_file):
+    raise ValueError(
+        "old_col_vocab_file and new_col_vocab_file must both be specified or "
+        "left unspecified. old_col_vocab_file='{}', new_col_vocab_file='{}'".
+        format(old_col_vocab_file, new_col_vocab_file))
+
+  remap_rows = new_row_vocab_file and old_row_vocab_file
+  remap_cols = new_col_vocab_file and old_col_vocab_file
+  if not (remap_rows or remap_cols):
+    raise ValueError(
+        "Must provide either row or column vocab files. If no remapping is "
+        "necessary, consider using `tf.contrib.framework.init_from_checkpoint` "
+        "instead.")
+
+  num_rows_present = num_rows_to_load
+  if remap_rows:
+    row_remapping, num_rows_present = (
+        gen_checkpoint_ops.generate_vocab_remapping(
+            new_vocab_file=new_row_vocab_file,
+            old_vocab_file=old_row_vocab_file,
+            new_vocab_offset=new_row_vocab_offset,
+            num_new_vocab=num_rows_to_load,
+            old_vocab_size=old_row_vocab_size))
+  else:
+    # Even when the rows are not being reordered, we still need to generate a
+    # remapping to account for initializing partitioned Variables (when
+    # new_row_vocab_offset is non-zero).
+    row_remapping = math_ops.range(
+        new_row_vocab_offset,
+        new_row_vocab_offset + num_rows_to_load,
+        dtype=dtypes.int64)
+
+  col_remapping = []
+  num_cols_present = new_col_vocab_size
+  if remap_cols:
+    col_remapping, num_cols_present = (
+        gen_checkpoint_ops.generate_vocab_remapping(
+            new_vocab_file=new_col_vocab_file,
+            old_vocab_file=old_col_vocab_file,
+            new_vocab_offset=0,  # Offset is unused for cols (no partitioning).
+            num_new_vocab=new_col_vocab_size))
+
+  init_vals = initializer([
+      num_rows_to_load * new_col_vocab_size -
+      num_rows_present * num_cols_present, 1
+  ])
+  return_tensor = gen_checkpoint_ops.load_and_remap_matrix(
+      ckpt_path=ckpt_path,
+      old_tensor_name=old_tensor_name,
+      row_remapping=row_remapping,
+      col_remapping=col_remapping,
+      initializing_values=init_vals,
+      num_rows=num_rows_to_load,
+      num_cols=new_col_vocab_size,
+      max_rows_in_memory=max_rows_in_memory)
+
+  # Add OOV row(s) and column(s).
+  if num_row_oov_buckets > 0:
+    init_row_oov_val = initializer([num_row_oov_buckets, new_col_vocab_size])
+    init_row_oov_val = ops.convert_to_tensor(init_row_oov_val)
+    return_tensor = array_ops.concat([return_tensor, init_row_oov_val], 0)
+  if num_col_oov_buckets > 0:
+    # We need to add any row OOV to the new column shape.
+    init_col_oov_val = initializer(
+        [num_rows_to_load + num_row_oov_buckets, num_col_oov_buckets])
+    init_col_oov_val = ops.convert_to_tensor(init_col_oov_val)
+    return_tensor = array_ops.concat([return_tensor, init_col_oov_val], 1)
+
+  return return_tensor
+
+
+def _load_and_remap_matrix_initializer(ckpt_path,
+                                       old_tensor_name,
+                                       new_row_vocab_size,
+                                       new_col_vocab_size,
+                                       old_row_vocab_size=-1,
+                                       old_row_vocab_file=None,
+                                       new_row_vocab_file=None,
+                                       old_col_vocab_file=None,
+                                       new_col_vocab_file=None,
+                                       num_row_oov_buckets=0,
+                                       num_col_oov_buckets=0,
+                                       initializer=None,
+                                       max_rows_in_memory=-1):
+  r"""Returns a var initializer for loading and remapping a 2-D (matrix) tensor.
+
+  The returned initializer loads a 2-D (matrix) `Tensor` with name
+  `old_tensor_name` from the checkpoint at `ckpt_path`. It will reorder the
+  rows/columns according to the specified vocab files and append additional
+  out-of-vocabulary rows/columns according to the number of OOV buckets.
+
+  The format of the file at the `{old,new}_{row,col}_vocab_file` path should be
+  a text file, with each line containing a single entity within the vocabulary.
+  Let the function `line_of(f, "x")` return the 0-indexed line number of the
+  entity "x" in file f, and the function `entity_at(f, i)` return the entity at
+  line i of file f. Then, row i of the new output matrix will be taken from row
+  `line_of(old_row_vocab_file, entity_at(new_row_vocab_file, i))` of the old
+  matrix. If any entity in `new_row_vocab_file` is not found in
+  `old_row_vocab_file`, that row is considered a "missing" row, and its values
+  will be initialized using the `initializer` arg. The same logic also applies
+  for the columns.
+
+  For example, assuming that:
+
+  * `old_row_vocab_file` contains "mercury\nvenus\nmars"
+  * `new_row_vocab_file` contains "venus\njupiter\nmercury"
+  * `old_col_vocab_file` contains "good\nbetter\nbest"
+  * `new_col_vocab_file` contains "good\nbest\nfantastic"
+  * `initializer` returns the natural numbers `[1, 2, 3, 4, ...]`
+  * `w(i, j)` represents the value from row i, column j of the old matrix
+
+  Then the new output matrix will look like:
+
+  `[[w(1, 0), w(1, 2), 1],
+    [2,       3,       4],
+    [w(0, 0), w(0, 2), 5]]`
+
+  If we further specify that:
+
+  * `num_row_oov_buckets` == 2
+  * `num_col_oov_buckets` == 1
+
+  Then the new output matrix will look like:
+
+  `[[w(1, 0), w(1, 2), 1,  12],
+    [2,       3,       4,  13],
+    [w(0, 0), w(0, 2), 5,  14],
+    [6,       7,       8,  15],
+    [9,       10,      11, 16]]`
+
+  If `{old,new}_row_vocab_file` are None, we assume that the old and new row
+  vocab files are the same, and no row remapping is done. If
+  `{old,new}_col_vocab_file` are None, we assume that the old and new column
+  vocab files are the same, and no column remapping is done.
+
+  The returned initializer only supports div-partitioning along the row axis. It
+  does not support partitioning along the column axis (as this is not common in
+  practice) or mod-partitioning.
+
+  NOTE: When this is used to warm-start variables, client code should use
+  `tf.lookup.index_table_from_tensor()` like
+  contrib/layers/python/layers/feature_column.py does, as opposed to
+  `tf.feature_to_id()` - in order to ensure the underlying lookup tables are the
+  same.
+
+  Args:
+    ckpt_path: Path to the TensorFlow checkpoint (version 2, `TensorBundle`)
+      from which the old matrix `Tensor` will be loaded.
+    old_tensor_name: Name of the 2-D `Tensor` to load from checkpoint.
+    new_row_vocab_size: `int` specifying the number of entries in
+      `new_row_vocab_file`. If no row remapping is needed (no row vocab
+      provided), this should be equal to the number of rows to load from the old
+      matrix (which can theoretically be smaller than the number of rows in the
+      old matrix).
+    new_col_vocab_size: `int` specifying the number of entries in
+      `new_col_vocab_file`. If no column remapping is needed (no column vocab
+      provided), this should be equal to the number of columns in the old
+      matrix.
+    old_row_vocab_size: The number of entries to consider in the old vocabulary.
+      With the default value of -1, the entire old row vocabulary file will be
+      used.  Otherwise, only the first `old_row_vocab_size` entries will be
+      considered for remapping.Must be smaller than the length of
+      `old_row_vocab_file`.  NOTE: we do not provide an equivalent
+      `old_col_vocab_size` for classes.
+    old_row_vocab_file: A scalar `Tensor` of type `string` containing the
+      path to the old row vocabulary file. Can be None, which represents no
+      remapping on the row axis.
+    new_row_vocab_file: A scalar `Tensor` of type `string` containing the path
+      to the new row vocabulary file. Can be None, which represents no remapping
+      on the row axis.
+    old_col_vocab_file: A scalar `Tensor` of type `string` containing the
+      path to the old column vocabulary file. Can be None, which represents no
+      remapping on the column axis.
+    new_col_vocab_file: A scalar `Tensor` of type `string` containing the path
+      to the new column vocabulary file. Can be None, which represents no
+      remapping on the column axis.
+    num_row_oov_buckets: `int` specifying the number of out-of-vocabulary rows
+      to append. Must be >= 0.
+    num_col_oov_buckets: `int` specifying the number of out-of-vocabulary
+      columns to append. Must be >= 0.
+    initializer: Initializer function to initialize missing values. Accepts a
+      1-D tensor as the arg to specify the shape of the returned tensor. If
+      `None`, defaults to using `zeros_initializer()`.
+    max_rows_in_memory: `int` specifying the maximum number of rows to load from
+      the checkpoint at once. If less than or equal to 0, the entire matrix will
+      be loaded into memory. Setting this arg trades increased disk reads for
+      lower memory usage.
+
+  Returns:
+    A variable initializer function that should be used to initialize a
+    (potentially partitioned) `Variable` whose complete shape is
+    `[new_row_vocab_size + num_row_oov_buckets, new_col_vocab_size +
+    num_col_oov_buckets]`.
+
+  Raises:
+    TypeError: If `initializer` is specified but not callable.
+  """
+  if initializer is None:
+    # TODO(b/25671353): Consider using sqrt(6/(fan_in + fan_out)) instead, from
+    # Glorot and Bengio, 2010.
+    initializer = init_ops.zeros_initializer()
+
+  if not callable(initializer):
+    raise TypeError(
+        "initializer must be callable, instead of being {} of type {}.".format(
+            initializer, type(initializer)))
+
+  def _initializer(shape, dtype=dtypes.float32, partition_info=None):
+    """Variable initializer.
+
+    Args:
+      shape: Shape of `Tensor` to return. Should include OOV on both axes.
+      dtype: Must be float32.
+      partition_info: variable_scope._PartitionInfo.
+
+    Returns:
+      `Tensor` of shape `shape`.
+
+    Raises:
+      TypeError: If `dtype` is anything other than float32.
+      ValueError: For shape mismatch upon invocation.
+    """
+    # Sanity checks.
+    if dtype != dtypes.float32:
+      raise TypeError(
+          "Currently, only float32 is supported. Received dtype: {}".format(
+              dtype))
+    if len(shape) != 2:
+      raise ValueError("Expected 2-dim shape, but received: {}".format(shape))
+    if shape[0] <= 0:
+      raise ValueError(
+          "Expected 1st dim of shape to be > 0, but received shape: {}".format(
+              shape))
+    if shape[1] != (new_col_vocab_size + num_col_oov_buckets):
+      raise ValueError(
+          "Expected 2nd dim of shape to be new_col_vocab_size ({}) + "
+          "num_col_oov_buckets ({}) = {}, but received shape: {}".format(
+              new_col_vocab_size, num_col_oov_buckets,
+              new_col_vocab_size + num_col_oov_buckets, shape))
+
+    offset = 0
+    if partition_info is not None:
+      offset = partition_info.single_offset(shape)
+
+    if offset + shape[0] > new_row_vocab_size + num_row_oov_buckets:
+      raise ValueError(
+          "Trying to initialize {} additional rows after {} rows have already "
+          "been initialized, which would exceed expected total row count of "
+          "new_row_vocab_size ({}) + num_row_oov_buckets ({}) = {}.".format(
+              shape[0], offset, new_row_vocab_size, num_row_oov_buckets,
+              new_row_vocab_size + num_row_oov_buckets))
+
+    row_oov_buckets_to_use = min(shape[0],
+                                 max(0, offset + shape[0] - new_row_vocab_size))
+    num_rows_to_load = shape[0] - row_oov_buckets_to_use
+
+    # We may be operating on an OOV-only partition, in which case we newly
+    # initialize all rows of this partition.
+    if offset > new_row_vocab_size:
+      if shape[0] != row_oov_buckets_to_use:
+        raise ValueError(
+            "Partitioned variable offset is greater than new vocab size and "
+            "not operating on OOV-only partition.")
+      return initializer(shape)
+
+    return _load_and_remap_matrix(
+        ckpt_path=ckpt_path,
+        old_tensor_name=old_tensor_name,
+        new_row_vocab_offset=offset,
+        num_rows_to_load=num_rows_to_load,
+        new_col_vocab_size=new_col_vocab_size,
+        initializer=initializer,
+        old_row_vocab_size=old_row_vocab_size,
+        old_row_vocab_file=old_row_vocab_file,
+        new_row_vocab_file=new_row_vocab_file,
+        old_col_vocab_file=old_col_vocab_file,
+        new_col_vocab_file=new_col_vocab_file,
+        num_row_oov_buckets=row_oov_buckets_to_use,
+        num_col_oov_buckets=num_col_oov_buckets,
+        max_rows_in_memory=max_rows_in_memory)
+
+  return _initializer
+
+
+def _load_embedding_initializer(ckpt_path,
+                                embedding_tensor_name,
+                                new_vocab_size,
+                                embedding_dim,
+                                old_vocab_file,
+                                new_vocab_file,
+                                old_vocab_size=-1,
+                                num_oov_buckets=0,
+                                initializer=None,
+                                max_rows_in_memory=-1):
+  """Returns a variable initializer for loading pre-trained embeddings.
+
+  Wrapper around `load_and_remap_matrix_initializer()` specialized for loading
+  embedding weights and remapping according to the provided vocab files. See
+  docs for `load_and_remap_matrix_initializer()` for more details.
+
+  NOTE: Only for use with div-partitioned variables / vocabularies.
+
+  Args:
+    ckpt_path: Path to the TensorFlow checkpoint (version 2, `TensorBundle`)
+      from which the old matrix `Tensor` will be loaded.
+    embedding_tensor_name: Name of the 2-D `Tensor` to load from checkpoint.
+    new_vocab_size: Number of entries in the new vocab.
+    embedding_dim: `int` specifying the dimension of the embedding vectors from
+      the checkpoint. Must match the number of columns in the old embedding
+      matrix.
+    old_vocab_file: A scalar `Tensor` of type `string` containing the
+      path to the old vocabulary file.
+    new_vocab_file: A scalar `Tensor` of type `string` containing the
+      path to the new vocabulary file.
+    old_vocab_size: The number of entries to consider in the old vocabulary.
+      With the default value of -1, the entire old row vocabulary file will be
+      used.  Otherwise, only the first `old_vocab_size` entries will be
+      considered for remapping.Must be smaller than the length of
+      `old_row_vocab_file`.
+    num_oov_buckets: `int` specifying the number of out-of-vocabulary
+      buckets to use. Must be >= 0.
+    initializer: Initializer function that accepts a 1-D tensor as the arg to
+      specify the shape of the returned tensor. If `None`, defaults to using
+      `truncated_normal_initializer()`.
+    max_rows_in_memory: `int` specifying the maximum number of rows to load from
+      the checkpoint at once. If less than or equal to 0, the entire matrix will
+      be loaded into memory. Setting this arg trades increased disk reads for
+      lower memory usage.
+
+  Returns:
+    A variable initializer function.
+  """
+  if initializer is None:
+    # TODO(b/25671353): This should be kept in sync with the stddev used by
+    # feature_column.py's _EmbeddingColumn.
+    initializer = init_ops.truncated_normal_initializer(
+        stddev=1.0 / math.sqrt(embedding_dim))
+
+  return _load_and_remap_matrix_initializer(
+      ckpt_path=ckpt_path,
+      old_tensor_name=embedding_tensor_name,
+      new_row_vocab_size=new_vocab_size,
+      new_col_vocab_size=embedding_dim,
+      old_row_vocab_size=old_vocab_size,
+      old_row_vocab_file=old_vocab_file,
+      new_row_vocab_file=new_vocab_file,
+      old_col_vocab_file=None,
+      new_col_vocab_file=None,
+      num_row_oov_buckets=num_oov_buckets,
+      num_col_oov_buckets=0,
+      initializer=initializer,
+      max_rows_in_memory=max_rows_in_memory)

tensorflow/python/training/checkpoint_ops_test.py

@@ -0,0 +1,383 @@
+# Copyright 2017 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 tests for Python wrappers around warm-starting."""
+import os
+
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import partitioned_variables
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import checkpoint_ops
+from tensorflow.python.training import saver as saver_lib
+
+
+@test_util.run_v1_only(
+    'This is to test V1 name-based checkpoints which is not supported in V2.')
+class LoadAndRemapWrappersTest(test.TestCase):
+  """Tests for the functionality of the Python wrappers."""
+
+  def setUp(self):
+    ops.reset_default_graph()
+    # Create the checkpoint file in a temporary directory.
+    checkpoint_prefix = os.path.join(self.get_temp_dir(), 'model')
+    # 0., 1., ..., 79. reshaped into [5, 16].
+    initializer = init_ops.constant_initializer(
+        np.reshape(np.linspace(0.0, 79, 5 * 16), (5, 16)))
+    with self.cached_session() as sess:
+      with variable_scope.variable_scope('some_scope'):
+        variable_scope.get_variable(name='embeddings', shape=[5, 16],
+                                    initializer=initializer)
+      self.evaluate(variables.global_variables_initializer())
+      saver = saver_lib.Saver()
+      saver.save(sess, checkpoint_prefix, global_step=5)
+    self.checkpoint_file = '{}-5'.format(checkpoint_prefix)
+
+    # Create the vocabulary files.
+    self.new_feature_vocab_file = os.path.join(
+        self.get_temp_dir(), 'new_feature_vocab.txt')
+    with open(self.new_feature_vocab_file, 'w') as f:
+      f.write('\n'.join(['zero', 'one', 'two', 'three', 'four']) + '\n')
+
+    self.old_feature_vocab_file = os.path.join(
+        self.get_temp_dir(), 'old_feature_vocab.txt')
+    with open(self.old_feature_vocab_file, 'w') as f:
+      f.write('\n'.join(['zero', 'one', 'two', 'three']) + '\n')
+
+    self.new_class_vocab_file = os.path.join(
+        self.get_temp_dir(), 'new_class_vocab.txt')
+    with open(self.new_class_vocab_file, 'w') as f:
+      f.write('\n'.join(['MISSING', 'knitting', 'flask', 'eminem']) + '\n')
+
+    self.old_class_vocab_file = os.path.join(
+        self.get_temp_dir(), 'old_class_vocab.txt')
+    with open(self.old_class_vocab_file, 'w') as f:
+      f.write('\n'.join(['knitting', 'eminem', 'MISSING']) + '\n')
+
+    self.init_val = 42
+
+    def _init_val_initializer(shape, dtype=None, partition_info=None):
+      del dtype, partition_info  # Unused by this unit-testing initializer.
+      return array_ops.tile(
+          constant_op.constant([[self.init_val]], dtype=dtypes.float32), shape)
+
+    self.initializer = _init_val_initializer
+
+  def test_load_and_remap_matrix(self):
+    """Tests the end-to-end loading / remapping of weights."""
+    # _load_and_remap_matrix() is the generalized wrapper that takes in row and
+    # column vocabulary files, calls the relevant remappings, and returns the
+    # weight matrix.  Take this example to be linear multi-class by providing
+    # both row and column vocabularies.
+    remapped_matrix = checkpoint_ops._load_and_remap_matrix(
+        new_row_vocab_file=self.new_feature_vocab_file,
+        old_row_vocab_file=self.old_feature_vocab_file,
+        num_rows_to_load=4,
+        new_col_vocab_file=self.new_class_vocab_file,
+        old_col_vocab_file=self.old_class_vocab_file,
+        new_col_vocab_size=4,
+        old_tensor_name='some_scope/embeddings',
+        ckpt_path=[self.checkpoint_file],
+        new_row_vocab_offset=1,
+        initializer=self.initializer,
+        num_row_oov_buckets=1,
+        num_col_oov_buckets=1)
+
+    # [4 in vocab + 1 oov features, 4 in vocab + 1 oov classes].  The offset
+    # means we read from the first line.
+    expected_remapped_matrix = np.concatenate(
+        [
+            np.reshape([18, 34, 50, self.init_val, self.init_val], [5, 1]),
+            np.reshape([16, 32, 48, self.init_val, self.init_val], [5, 1]),
+            np.reshape([self.init_val] * 5, [5, 1]),
+            np.reshape([17, 33, 49, self.init_val, self.init_val], [5, 1]),
+            np.reshape([self.init_val] * 5, [5, 1])
+        ],
+        axis=1)
+
+    with self.cached_session():
+      self.assertAllClose(expected_remapped_matrix,
+                          self.evaluate(remapped_matrix))
+
+  def test_load_and_remap_output_layer_weight_initializer_linear(self):
+    """Tests for the output layer initializer in the linear multi-class case."""
+    loading_initializer = (checkpoint_ops._load_and_remap_matrix_initializer(
+        new_row_vocab_size=5,
+        new_col_vocab_file=self.new_class_vocab_file,
+        old_col_vocab_file=self.old_class_vocab_file,
+        new_col_vocab_size=4,
+        old_tensor_name='some_scope/embeddings',
+        ckpt_path=[self.checkpoint_file],
+        new_row_vocab_file=self.new_feature_vocab_file,
+        old_row_vocab_file=self.old_feature_vocab_file,
+        num_row_oov_buckets=1,
+        num_col_oov_buckets=1,
+        initializer=self.initializer))
+
+    # The new weight matrix is of size
+    # [5 feature vocab + 1 feature OOV, 4 class vocab + 1 class OOV].  Use a
+    # partitioned variable to confirm that the offset logic works.
+    expected_remapped_matrix = np.concatenate(
+        [
+            np.reshape([2, 18, 34, 50, self.init_val, self.init_val], [6, 1]),
+            np.reshape([0, 16, 32, 48, self.init_val, self.init_val], [6, 1]),
+            np.reshape([self.init_val] * 6, [6, 1]),
+            np.reshape([1, 17, 33, 49, self.init_val, self.init_val], [6, 1]),
+            np.reshape([self.init_val] * 6, [6, 1])
+        ],
+        axis=1)
+    remapped_matrix = variable_scope.get_variable(
+        name='linear/obtained_weight_matrix',
+        shape=[6, 5],
+        initializer=loading_initializer,
+        partitioner=partitioned_variables.fixed_size_partitioner(2))
+
+    with self.cached_session():
+      self.evaluate(variables.global_variables_initializer())
+      self.assertAllClose(expected_remapped_matrix, remapped_matrix.as_tensor())
+
+  def test_load_and_remap_output_layer_weight_initializer_dnn_output(self):
+    """Tests for the output layer initializer in the DNN output case."""
+    loading_initializer = (checkpoint_ops._load_and_remap_matrix_initializer(
+        new_row_vocab_size=5,
+        new_col_vocab_file=self.new_class_vocab_file,
+        old_col_vocab_file=self.old_class_vocab_file,
+        new_col_vocab_size=4,
+        old_tensor_name='some_scope/embeddings',
+        ckpt_path=[self.checkpoint_file],
+        num_col_oov_buckets=1,
+        initializer=self.initializer))
+
+    # The new weight matrix is of size
+    # [5-sized input layer, 4 class vocab + 1 class OOV].
+    expected_remapped_matrix = np.concatenate(
+        [
+            np.reshape([2, 18, 34, 50, 66], [5, 1]),
+            np.reshape([0, 16, 32, 48, 64], [5, 1]),
+            np.reshape([self.init_val] * 5, [5, 1]),
+            np.reshape([1, 17, 33, 49, 65], [5, 1]),
+            np.reshape([self.init_val] * 5, [5, 1])
+        ],
+        axis=1)
+    remapped_matrix = variable_scope.get_variable(
+        name='dnn_output/obtained_weight_matrix',
+        shape=[5, 5],
+        initializer=loading_initializer,
+        partitioner=partitioned_variables.fixed_size_partitioner(2))
+
+    with self.cached_session():
+      self.evaluate(variables.global_variables_initializer())
+      self.assertAllClose(expected_remapped_matrix, remapped_matrix.as_tensor())
+
+  def test_initializer_with_oov_only_partition(self):
+    """Tests for the output layer initializer where one partition is all OOV."""
+    loading_initializer = (checkpoint_ops._load_and_remap_matrix_initializer(
+        new_row_vocab_size=5,
+        new_col_vocab_file=self.new_class_vocab_file,
+        old_col_vocab_file=self.old_class_vocab_file,
+        new_col_vocab_size=4,
+        old_tensor_name='some_scope/embeddings',
+        ckpt_path=[self.checkpoint_file],
+        new_row_vocab_file=self.new_feature_vocab_file,
+        old_row_vocab_file=self.old_feature_vocab_file,
+        num_row_oov_buckets=5,
+        num_col_oov_buckets=1,
+        initializer=self.initializer))
+
+    # The new weight matrix is of size
+    # [5 feature vocab + 5 feature OOV, 4 class vocab + 1 class OOV].  The
+    # second partition has only OOV.
+    expected_remapped_matrix = np.concatenate(
+        [
+            np.reshape([2, 18, 34, 50] + [self.init_val] * 6, [10, 1]),
+            np.reshape([0, 16, 32, 48] + [self.init_val] * 6, [10, 1]),
+            np.reshape([self.init_val] * 10, [10, 1]),
+            np.reshape([1, 17, 33, 49] + [self.init_val] * 6, [10, 1]),
+            np.reshape([self.init_val] * 10, [10, 1]),
+        ],
+        axis=1)
+    remapped_matrix = variable_scope.get_variable(
+        name='linear_all_oov/obtained_weight_matrix',
+        shape=[10, 5],
+        initializer=loading_initializer,
+        partitioner=partitioned_variables.fixed_size_partitioner(2))
+
+    with self.cached_session():
+      self.evaluate(variables.global_variables_initializer())
+      self.assertAllClose(expected_remapped_matrix, remapped_matrix.as_tensor())
+
+  def test_load_and_remap_linear_multiclass_initializer_default_init(self):
+    """Tests where the zeros_initializer default is used for linear."""
+    loading_initializer = (checkpoint_ops._load_and_remap_matrix_initializer(
+        new_row_vocab_size=5,
+        new_col_vocab_file=self.new_class_vocab_file,
+        old_col_vocab_file=self.old_class_vocab_file,
+        new_col_vocab_size=4,
+        old_tensor_name='some_scope/embeddings',
+        ckpt_path=[self.checkpoint_file],
+        new_row_vocab_file=self.new_feature_vocab_file,
+        old_row_vocab_file=self.old_feature_vocab_file,
+        num_row_oov_buckets=1,
+        num_col_oov_buckets=1))
+
+    # Same as test_initializer_with_oov_only_partition, but with zero
+    # initialization.
+    expected_remapped_matrix = np.concatenate(
+        [
+            np.reshape([2, 18, 34, 50, 0, 0], [6, 1]),
+            np.reshape([0, 16, 32, 48, 0, 0], [6, 1]),
+            np.reshape([0] * 6, [6, 1]),
+            np.reshape([1, 17, 33, 49, 0, 0], [6, 1]),
+            np.reshape([0] * 6, [6, 1])
+        ],
+        axis=1)
+    remapped_matrix = variable_scope.get_variable(
+        name='linear_init_fallback/obtained_weight_matrix',
+        shape=[6, 5],
+        initializer=loading_initializer,
+        partitioner=partitioned_variables.fixed_size_partitioner(2))
+
+    with self.cached_session():
+      self.evaluate(variables.global_variables_initializer())
+      self.assertAllClose(expected_remapped_matrix, remapped_matrix.as_tensor())
+
+  def test_load_embedding_initializer(self):
+    """Tests for the load_embedding_initializer wrapper."""
+    embedding_loading_initializer = (checkpoint_ops._load_embedding_initializer(
+        new_vocab_file=self.new_feature_vocab_file,
+        old_vocab_file=self.old_feature_vocab_file,
+        new_vocab_size=5,
+        embedding_dim=16,
+        embedding_tensor_name='some_scope/embeddings',
+        ckpt_path=[self.checkpoint_file],
+        num_oov_buckets=1,
+        initializer=self.initializer))
+
+    # The new weight matrix is of size
+    # [5 feature vocab + 1 feature OOV, 16 (embedding dimension)], where the
+    # last vocab row (2nd last row) is newly initialized (wasn't found in
+    # previous vocab) and the actual last row is OOV and also newly initialized.
+    # Use a partitioned variable to confirm that the offset logic works.
+    expected_remapped_embeddings = np.concatenate(
+        [
+            np.reshape(range(64), [4, 16]),
+            np.reshape([self.init_val] * 32, [2, 16]),
+        ],
+        axis=0)
+    remapped_embeddings = variable_scope.get_variable(
+        name='embedding/obtained_embedding_matrix',
+        shape=[6, 16],
+        initializer=embedding_loading_initializer,
+        partitioner=partitioned_variables.fixed_size_partitioner(2))
+
+    with self.cached_session():
+      self.evaluate(variables.global_variables_initializer())
+      self.assertAllClose(expected_remapped_embeddings,
+                          remapped_embeddings.as_tensor())
+
+  def test_load_embedding_initializer_large_oov(self):
+    """Tests for the large OOV case for load_embedding_initializer wrapper."""
+    self.new_feature_vocab_file = os.path.join(
+        self.get_temp_dir(), 'new_feature_vocab.txt')
+    with open(self.new_feature_vocab_file, 'w') as f:
+      f.write('\n'.join(['one', 'zero', 'two', 'four']) + '\n')
+
+    # Checkpoint has 5 entries, 3 of which correspond to OOV.
+    self.old_feature_vocab_file = os.path.join(
+        self.get_temp_dir(), 'old_feature_vocab.txt')
+    with open(self.old_feature_vocab_file, 'w') as f:
+      f.write('\n'.join(['zero', 'one']) + '\n')
+
+    embedding_loading_initializer = (checkpoint_ops._load_embedding_initializer(
+        new_vocab_file=self.new_feature_vocab_file,
+        old_vocab_file=self.old_feature_vocab_file,
+        new_vocab_size=4,
+        embedding_dim=16,
+        embedding_tensor_name='some_scope/embeddings',
+        ckpt_path=[self.checkpoint_file],
+        num_oov_buckets=5,
+        initializer=self.initializer))
+
+    # The new weight matrix is of size
+    # [4 feature vocab + 5 feature OOV, 16 (embedding dimension)], where the
+    # 3rd and 4th rows are not found in the old vocabulary and therefore newly
+    # initialized.  The last five rows are OOV and also newly initialized.
+    # Use a partitioned variable to confirm that the offset logic works.
+    expected_remapped_embeddings = np.concatenate(
+        [
+            np.reshape(range(16, 32), [1, 16]),
+            np.reshape(range(16), [1, 16]),
+            np.reshape([self.init_val] * 112, [7, 16]),
+        ],
+        axis=0)
+    remapped_embeddings = variable_scope.get_variable(
+        name='embedding/obtained_embedding_matrix',
+        shape=[9, 16],
+        initializer=embedding_loading_initializer,
+        partitioner=partitioned_variables.fixed_size_partitioner(2))
+
+    with self.cached_session():
+      self.evaluate(variables.global_variables_initializer())
+      self.assertAllClose(expected_remapped_embeddings,
+                          remapped_embeddings.as_tensor())
+
+  def test_load_embedding_initializer_old_row_vocab(self):
+    """Tests for load_embedding_initializer where we constrain old vocab."""
+    embedding_loading_initializer = (
+        checkpoint_ops._load_embedding_initializer(
+            new_vocab_file=self.new_feature_vocab_file,
+            old_vocab_file=self.old_feature_vocab_file,
+            # Considered old vocabulary becomes ['zero', 'one', 'two'].  This
+            # means 'three' in the new vocabulary is newly initialized.
+            old_vocab_size=3,
+            new_vocab_size=5,
+            embedding_dim=16,
+            embedding_tensor_name='some_scope/embeddings',
+            ckpt_path=[self.checkpoint_file],
+            num_oov_buckets=1,
+            initializer=self.initializer))
+
+    # The new weight matrix is of size
+    # [5 feature vocab + 1 feature OOV, 16 (embedding dimension)], where the
+    # last vocab row (2nd last row) is newly initialized (wasn't found in
+    # previous vocab) and the actual last row is OOV and also newly initialized.
+    # Use a partitioned variable to confirm that the offset logic works.
+    expected_remapped_embeddings = np.concatenate(
+        [
+            np.reshape(range(48), [3, 16]),
+            np.reshape([self.init_val] * 48, [3, 16]),
+        ],
+        axis=0)
+    remapped_embeddings = variable_scope.get_variable(
+        name='embedding/obtained_embedding_matrix',
+        shape=[6, 16],
+        initializer=embedding_loading_initializer,
+        partitioner=partitioned_variables.fixed_size_partitioner(2))
+
+    with self.cached_session():
+      self.evaluate(variables.global_variables_initializer())
+      self.assertAllClose(expected_remapped_embeddings,
+                          remapped_embeddings.as_tensor())
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/training/checkpoint_state.proto

@@ -0,0 +1,38 @@
+// Copyright 2026 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.
+// ==============================================================================
+
+syntax = "proto3";
+
+package tensorflow;
+
+option cc_enable_arenas = true;
+
+// Protocol buffer representing the checkpoint state.
+message CheckpointState {
+  // Path to the most-recent model checkpoint.
+  string model_checkpoint_path = 1;
+
+  // Paths to all not-yet-deleted model checkpoints, sorted from oldest to
+  // newest.
+  // Note that the value of model_checkpoint_path should be the last item in
+  // this list.
+  repeated string all_model_checkpoint_paths = 2;
+  // Unix timestamps corresponding to all_model_checkpoint_paths, indicating
+  // when each checkpoint was created.
+  repeated double all_model_checkpoint_timestamps = 3;
+  // Unix timestamp indicating the creation time for the last preserved
+  // checkpoint.
+  double last_preserved_timestamp = 4;
+}

tensorflow/python/training/checkpoint_utils.py

@@ -0,0 +1,570 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Tools to work with name-based checkpoints.
+
+While some of these symbols also work with the TF2 object-based checkpoints,
+they are not recommended for TF2. Please check  `tensorflow/python/checkpoint`
+for newer utilities built to work with TF2 checkpoints.
+"""
+
+from collections import abc
+import os
+import time
+
+from tensorflow.python.checkpoint import checkpoint_management
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import io_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variable_scope as vs
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import gfile
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.training import py_checkpoint_reader
+from tensorflow.python.training.saving import saveable_object_util
+from tensorflow.python.util.tf_export import tf_export
+
+
+__all__ = [
+    "load_checkpoint", "load_variable", "list_variables",
+    "checkpoints_iterator", "init_from_checkpoint"
+]
+
+
+@tf_export("train.load_checkpoint")
+def load_checkpoint(ckpt_dir_or_file):
+  """Returns `CheckpointReader` for checkpoint found in `ckpt_dir_or_file`.
+
+  If `ckpt_dir_or_file` resolves to a directory with multiple checkpoints,
+  reader for the latest checkpoint is returned.
+
+  Example usage:
+
+  ```python
+  import tensorflow as tf
+  a = tf.Variable(1.0)
+  b = tf.Variable(2.0)
+  ckpt = tf.train.Checkpoint(var_list={'a': a, 'b': b})
+  ckpt_path = ckpt.save('tmp-ckpt')
+  reader= tf.train.load_checkpoint(ckpt_path)
+  print(reader.get_tensor('var_list/a/.ATTRIBUTES/VARIABLE_VALUE'))  # 1.0
+  ```
+
+  Args:
+    ckpt_dir_or_file: Directory with checkpoints file or path to checkpoint
+      file.
+
+  Returns:
+    `CheckpointReader` object.
+
+  Raises:
+    ValueError: If `ckpt_dir_or_file` resolves to a directory with no
+      checkpoints.
+  """
+  filename = _get_checkpoint_filename(ckpt_dir_or_file)
+  if filename is None:
+    raise ValueError("Couldn't find 'checkpoint' file or checkpoints in "
+                     "given directory %s" % ckpt_dir_or_file)
+  return py_checkpoint_reader.NewCheckpointReader(filename)
+
+
+@tf_export("train.load_variable")
+def load_variable(ckpt_dir_or_file, name):
+  """Returns the tensor value of the given variable in the checkpoint.
+
+  When the variable name is unknown, you can use `tf.train.list_variables` to
+  inspect all the variable names.
+
+  Example usage:
+
+  ```python
+  import tensorflow as tf
+  a = tf.Variable(1.0)
+  b = tf.Variable(2.0)
+  ckpt = tf.train.Checkpoint(var_list={'a': a, 'b': b})
+  ckpt_path = ckpt.save('tmp-ckpt')
+  var= tf.train.load_variable(
+      ckpt_path, 'var_list/a/.ATTRIBUTES/VARIABLE_VALUE')
+  print(var)  # 1.0
+  ```
+
+  Args:
+    ckpt_dir_or_file: Directory with checkpoints file or path to checkpoint.
+    name: Name of the variable to return.
+
+  Returns:
+    A numpy `ndarray` with a copy of the value of this variable.
+  """
+  # TODO(b/29227106): Fix this in the right place and remove this.
+  if name.endswith(":0"):
+    name = name[:-2]
+  reader = load_checkpoint(ckpt_dir_or_file)
+  return reader.get_tensor(name)
+
+
+@tf_export("train.list_variables")
+def list_variables(ckpt_dir_or_file):
+  """Lists the checkpoint keys and shapes of variables in a checkpoint.
+
+  Checkpoint keys are paths in a checkpoint graph.
+
+  Example usage:
+
+  ```python
+  import tensorflow as tf
+  import os
+  ckpt_directory = "/tmp/training_checkpoints/ckpt"
+  ckpt = tf.train.Checkpoint(optimizer=optimizer, model=model)
+  manager = tf.train.CheckpointManager(ckpt, ckpt_directory, max_to_keep=3)
+  train_and_checkpoint(model, manager)
+  tf.train.list_variables(manager.latest_checkpoint)
+  ```
+
+  Args:
+    ckpt_dir_or_file: Directory with checkpoints file or path to checkpoint.
+
+  Returns:
+    List of tuples `(key, shape)`.
+  """
+  reader = load_checkpoint(ckpt_dir_or_file)
+  variable_map = reader.get_variable_to_shape_map()
+  names = sorted(variable_map.keys())
+  result = []
+  for name in names:
+    result.append((name, variable_map[name]))
+  return result
+
+
+def wait_for_new_checkpoint(checkpoint_dir,
+                            last_checkpoint=None,
+                            seconds_to_sleep=1,
+                            timeout=None):
+  """Waits until a new checkpoint file is found.
+
+  Args:
+    checkpoint_dir: The directory in which checkpoints are saved.
+    last_checkpoint: The last checkpoint path used or `None` if we're expecting
+      a checkpoint for the first time.
+    seconds_to_sleep: The number of seconds to sleep for before looking for a
+      new checkpoint.
+    timeout: The maximum number of seconds to wait. If left as `None`, then the
+      process will wait indefinitely.
+
+  Returns:
+    a new checkpoint path, or None if the timeout was reached.
+  """
+  logging.info("Waiting for new checkpoint at %s", checkpoint_dir)
+  stop_time = time.time() + timeout if timeout is not None else None
+  while True:
+    checkpoint_path = checkpoint_management.latest_checkpoint(checkpoint_dir)
+    if checkpoint_path is None or checkpoint_path == last_checkpoint:
+      if stop_time is not None and time.time() + seconds_to_sleep > stop_time:
+        return None
+      time.sleep(seconds_to_sleep)
+    else:
+      logging.info("Found new checkpoint at %s", checkpoint_path)
+      return checkpoint_path
+
+
+@tf_export("train.checkpoints_iterator")
+def checkpoints_iterator(checkpoint_dir,
+                         min_interval_secs=0,
+                         timeout=None,
+                         timeout_fn=None):
+  """Continuously yield new checkpoint files as they appear.
+
+  The iterator only checks for new checkpoints when control flow has been
+  reverted to it. This means it can miss checkpoints if your code takes longer
+  to run between iterations than `min_interval_secs` or the interval at which
+  new checkpoints are written.
+
+  The `timeout` argument is the maximum number of seconds to block waiting for
+  a new checkpoint.  It is used in combination with the `timeout_fn` as
+  follows:
+
+  * If the timeout expires and no `timeout_fn` was specified, the iterator
+    stops yielding.
+  * If a `timeout_fn` was specified, that function is called and if it returns
+    a true boolean value the iterator stops yielding.
+  * If the function returns a false boolean value then the iterator resumes the
+    wait for new checkpoints.  At this point the timeout logic applies again.
+
+  This behavior gives control to callers on what to do if checkpoints do not
+  come fast enough or stop being generated.  For example, if callers have a way
+  to detect that the training has stopped and know that no new checkpoints
+  will be generated, they can provide a `timeout_fn` that returns `True` when
+  the training has stopped.  If they know that the training is still going on
+  they return `False` instead.
+
+  Args:
+    checkpoint_dir: The directory in which checkpoints are saved.
+    min_interval_secs: The minimum number of seconds between yielding
+      checkpoints.
+    timeout: The maximum number of seconds to wait between checkpoints. If left
+      as `None`, then the process will wait indefinitely.
+    timeout_fn: Optional function to call after a timeout.  If the function
+      returns True, then it means that no new checkpoints will be generated and
+      the iterator will exit.  The function is called with no arguments.
+
+  Yields:
+    String paths to latest checkpoint files as they arrive.
+  """
+  checkpoint_path = None
+  while True:
+    new_checkpoint_path = wait_for_new_checkpoint(
+        checkpoint_dir, checkpoint_path, timeout=timeout)
+    if new_checkpoint_path:
+      start = time.time()
+      checkpoint_path = new_checkpoint_path
+      yield checkpoint_path
+      time_to_next_eval = start + min_interval_secs - time.time()
+      if time_to_next_eval > 0:
+        time.sleep(time_to_next_eval)
+    elif timeout_fn is None:
+      # No way to determine if any additional checkpoints will be generated.
+      logging.info("Timed-out waiting for a checkpoint (without timeout_fn).")
+      return
+    # Check if more new checkpoints may be generated.
+    if timeout_fn is not None and timeout_fn():
+      # The timeout_fn indicated that we are truly done.
+      logging.info("Done. No new checkpoints will be generated.")
+      return
+
+
+@tf_export(v1=["train.init_from_checkpoint"])
+def init_from_checkpoint(ckpt_dir_or_file, assignment_map):
+  """Replaces `tf.Variable` initializers so they load from a checkpoint file.
+
+  @compatibility(TF2)
+  `tf.compat.v1.train.init_from_checkpoint` is not recommended for restoring
+  variable values in TF2.
+
+  To restore checkpoints in TF2, please use
+  `tf.keras.Model.load_weights` or `tf.train.Checkpoint.restore`. These APIs use
+  use an [object-based method of checkpointing]
+  (https://www.tensorflow.org/guide/checkpoint#loading_mechanics), while
+  `tf.compat.v1.init_from_checkpoint` relies on a more-fragile variable-name
+  based method of checkpointing. There is no object-based equivalent of
+  `init_from_checkpoint` in TF2.
+
+  Please re-write your checkpoints immediately using the object-based APIs,
+  see [migration guide]
+  (https://www.tensorflow.org/guide/migrate#checkpoint_compatibility) for more
+  details.
+
+  You can load a name-based checkpoint written by `tf.compat.v1.train.Saver`
+  using `tf.train.Checkpoint.restore` or `tf.keras.Model.load_weights`. However,
+  you may have to change the names of the variables in your model to match the
+  variable names in the name-based checkpoint, which can be viewed with
+  `tf.train.list_variables(path)`.
+
+  Another option is to create an `assignment_map` that maps the name of the
+  variables in the name-based checkpoint to the variables in your model, eg:
+  ```
+  {
+      'sequential/dense/bias': model.variables[0],
+      'sequential/dense/kernel': model.variables[1]
+  }
+  ```
+  and use `tf.compat.v1.train.init_from_checkpoint(path, assignment_map)` to
+  restore the name-based checkpoint.
+
+  After restoring, re-encode your checkpoint using `tf.train.Checkpoint.save`
+  or `tf.keras.Model.save_weights`.
+
+  @end_compatibility
+
+  Values are not loaded immediately, but when the initializer is run
+  (typically by running a `tf.compat.v1.global_variables_initializer` op).
+
+  Note: This overrides default initialization ops of specified variables and
+  redefines dtype.
+
+  Assignment map supports following syntax:
+
+  * `'checkpoint_scope_name/': 'scope_name/'` - will load all variables in
+    current `scope_name` from `checkpoint_scope_name` with matching tensor
+    names.
+  * `'checkpoint_scope_name/some_other_variable': 'scope_name/variable_name'` -
+    will initialize `scope_name/variable_name` variable
+    from `checkpoint_scope_name/some_other_variable`.
+  * `'scope_variable_name': variable` - will initialize given `tf.Variable`
+    object with tensor 'scope_variable_name' from the checkpoint.
+  * `'scope_variable_name': list(variable)` - will initialize list of
+    partitioned variables with tensor 'scope_variable_name' from the checkpoint.
+  * `'/': 'scope_name/'` - will load all variables in current `scope_name` from
+    checkpoint's root (e.g. no scope).
+
+  Supports loading into partitioned variables, which are represented as
+  `'<variable>/part_<part #>'`.
+
+  Assignment map can be a dict, or a list of pairs.  The latter is
+  necessary to initialize multiple variables in the current graph from
+  the same variable in the checkpoint.
+
+  Example:
+
+  ```python
+
+  # Say, '/tmp/model.ckpt' has the following tensors:
+  #  -- name='old_scope_1/var1', shape=[20, 2]
+  #  -- name='old_scope_1/var2', shape=[50, 4]
+  #  -- name='old_scope_2/var3', shape=[100, 100]
+
+  # Create new model's variables
+  with tf.compat.v1.variable_scope('new_scope_1'):
+    var1 = tf.compat.v1.get_variable('var1', shape=[20, 2],
+                           initializer=tf.compat.v1.zeros_initializer())
+  with tf.compat.v1.variable_scope('new_scope_2'):
+    var2 = tf.compat.v1.get_variable('var2', shape=[50, 4],
+                           initializer=tf.compat.v1.zeros_initializer())
+    # Partition into 5 variables along the first axis.
+    var3 = tf.compat.v1.get_variable(name='var3', shape=[100, 100],
+                           initializer=tf.compat.v1.zeros_initializer(),
+                           partitioner=lambda shape, dtype: [5, 1])
+
+  # Initialize all variables in `new_scope_1` from `old_scope_1`.
+  init_from_checkpoint('/tmp/model.ckpt', {'old_scope_1/': 'new_scope_1/'})
+
+  # Use names to specify which variables to initialize from checkpoint.
+  init_from_checkpoint('/tmp/model.ckpt',
+                       {'old_scope_1/var1': 'new_scope_1/var1',
+                        'old_scope_1/var2': 'new_scope_2/var2'})
+
+  # Or use tf.Variable objects to identify what to initialize.
+  init_from_checkpoint('/tmp/model.ckpt',
+                       {'old_scope_1/var1': var1,
+                        'old_scope_1/var2': var2})
+
+  # Initialize partitioned variables using variable's name
+  init_from_checkpoint('/tmp/model.ckpt',
+                       {'old_scope_2/var3': 'new_scope_2/var3'})
+
+  # Or specify the list of tf.Variable objects.
+  init_from_checkpoint('/tmp/model.ckpt',
+                       {'old_scope_2/var3': var3._get_variable_list()})
+
+  ```
+
+  Args:
+    ckpt_dir_or_file: Directory with checkpoints file or path to checkpoint.
+    assignment_map: Dict, or a list of key-value pairs, where keys are names
+      of the variables in the checkpoint and values are current variables or
+      names of current variables (in default graph).
+
+  Raises:
+    ValueError: If missing variables in current graph, or if missing
+      checkpoints or tensors in checkpoints.
+
+  """
+  init_from_checkpoint_fn = lambda _: _init_from_checkpoint(
+      ckpt_dir_or_file, assignment_map)
+  if distribute_lib.get_cross_replica_context():
+    init_from_checkpoint_fn(None)
+  else:
+    distribute_lib.get_replica_context().merge_call(
+        init_from_checkpoint_fn)
+
+
+def _init_from_checkpoint(ckpt_dir_or_file, assignment_map):
+  """See `init_from_checkpoint` for documentation."""
+  ckpt_file = _get_checkpoint_filename(ckpt_dir_or_file)
+  reader = load_checkpoint(ckpt_dir_or_file)
+  variable_map = reader.get_variable_to_shape_map()
+  if isinstance(assignment_map, abc.Mapping):
+    assignment_map = assignment_map.items()
+
+  # We only want to sort by tensor names.
+  sort_key = lambda pair: pair[0]
+
+  for tensor_name_in_ckpt, current_var_or_name in sorted(
+      assignment_map, key=sort_key):
+    var = None
+    # Check if this is Variable object or list of Variable objects (in case of
+    # partitioned variables).
+    if _is_variable(current_var_or_name) or (
+        isinstance(current_var_or_name, list)
+        and all(_is_variable(v) for v in current_var_or_name)):
+      var = current_var_or_name
+    else:
+      store_vars = vs._get_default_variable_store()._vars  # pylint:disable=protected-access
+      # Check if this variable is in var_store.
+      var = store_vars.get(current_var_or_name, None)
+      # Also check if variable is partitioned as list.
+      if var is None:
+        var = _collect_partitioned_variable(current_var_or_name, store_vars)
+    if var is not None:
+      # If 1 to 1 mapping was provided, find variable in the checkpoint.
+      if tensor_name_in_ckpt not in variable_map:
+        raise ValueError("Tensor %s is not found in %s checkpoint %s" % (
+            tensor_name_in_ckpt, ckpt_dir_or_file, variable_map
+        ))
+      if _is_variable(var):
+        # Additional at-call-time checks.
+        if not var.get_shape().is_compatible_with(
+            variable_map[tensor_name_in_ckpt]):
+          raise ValueError(
+              "Shape of variable %s (%s) doesn't match with shape of "
+              "tensor %s (%s) from checkpoint reader." % (
+                  var.name, str(var.get_shape()),
+                  tensor_name_in_ckpt, str(variable_map[tensor_name_in_ckpt])
+              ))
+        var_name = var.name
+      else:
+        var_name = ",".join(v.name for v in var)
+      _set_variable_or_list_initializer(var, ckpt_file, tensor_name_in_ckpt)
+      logging.debug("Initialize variable %s from checkpoint %s with %s",
+                    var_name, ckpt_dir_or_file, tensor_name_in_ckpt)
+    else:
+      scopes = ""
+      # TODO(vihanjain): Support list of 'current_var_or_name' here.
+      if "/" in current_var_or_name:
+        scopes = current_var_or_name[:current_var_or_name.rindex("/")]
+      if not tensor_name_in_ckpt.endswith("/"):
+        raise ValueError(
+            "Assignment map with scope only name {} should map to scope only "
+            "{}. Should be 'scope/': 'other_scope/'.".format(
+                scopes, tensor_name_in_ckpt))
+      # If scope to scope mapping was provided, find all variables in the scope
+      # and create variable to variable mapping.
+      scope_variables = set()
+      for var_name in store_vars:
+        if not scopes or var_name.startswith(scopes + "/"):
+          # Consume /part_ if partitioned variable.
+          if "/part_" in var_name:
+            var_name = var_name[:var_name.index("/part_")]
+          scope_variables.add(var_name)
+      for var_name in sorted(scope_variables):
+        # Lookup name with specified prefix and suffix from current variable.
+        # If tensor_name given is '/' (root), don't use it for full name.
+        full_tensor_name = var_name[len(scopes):]
+        if current_var_or_name != "/":
+          full_tensor_name = full_tensor_name[1:]
+        if tensor_name_in_ckpt != "/":
+          full_tensor_name = tensor_name_in_ckpt + full_tensor_name
+        # Remove trailing '/', if any, in the full_tensor_name
+        if full_tensor_name.endswith("/"):
+          full_tensor_name = full_tensor_name[:-1]
+        if full_tensor_name not in variable_map:
+          raise ValueError(
+              "Tensor %s (%s in %s) is not found in %s checkpoint" % (
+                  full_tensor_name, var_name[len(scopes) + 1:],
+                  tensor_name_in_ckpt, ckpt_dir_or_file
+              ))
+        var = store_vars.get(var_name, None)
+        if var is None:
+          var = _collect_partitioned_variable(var_name, store_vars)
+        _set_variable_or_list_initializer(var, ckpt_file, full_tensor_name)
+        logging.debug("Initialize variable %s from checkpoint %s with %s",
+                      var_name, ckpt_dir_or_file, full_tensor_name)
+
+
+def _get_checkpoint_filename(ckpt_dir_or_file):
+  """Returns checkpoint filename given directory or specific checkpoint file."""
+  if isinstance(ckpt_dir_or_file, os.PathLike):
+    ckpt_dir_or_file = os.fspath(ckpt_dir_or_file)
+  if gfile.IsDirectory(ckpt_dir_or_file):
+    return checkpoint_management.latest_checkpoint(ckpt_dir_or_file)
+  return ckpt_dir_or_file
+
+
+def _set_checkpoint_initializer(variable,
+                                ckpt_file,
+                                tensor_name,
+                                slice_spec,
+                                name="checkpoint_initializer"):
+  """Overrides given variable's initialization op.
+
+  Sets variable initializer to assign op that initializes variable from tensor's
+  value in the checkpoint.
+
+  Args:
+    variable: `tf.Variable` object.
+    ckpt_file: string, full path of the checkpoint.
+    tensor_name: Name of the tensor to load from the checkpoint.
+    slice_spec: Slice specification for loading partitioned tensors.
+    name: Name of the operation.
+  """
+  base_type = variable.dtype.base_dtype
+  # Do not colocate with variable since RestoreV2 op only runs on CPU and
+  # colocation will force variable (and other ops that colocate with variable)
+  # to be on CPU as well. It is okay to place the variable's initializer op on
+  # CPU since it will only be run once at the start.
+  with ops.device(variable.device), ops.device("/cpu:0"):
+    restore_op = io_ops.restore_v2(
+        ckpt_file, [tensor_name], [slice_spec], [base_type], name=name)[0]
+
+    names_to_saveables = saveable_object_util.op_list_to_dict([variable])
+    saveable_objects = []
+    for name, op in names_to_saveables.items():
+      for s in saveable_object_util.saveable_objects_for_op(op, name):
+        saveable_objects.append(s)
+
+    assert len(saveable_objects) == 1  # Should be only one variable.
+  init_op = saveable_objects[0].restore([restore_op], restored_shapes=None)
+
+  # pylint:disable=protected-access
+  variable._initializer_op = init_op
+  restore_op.set_shape(variable.shape)
+  variable._initial_value = restore_op
+  # pylint:enable=protected-access
+
+
+def _set_variable_or_list_initializer(variable_or_list, ckpt_file,
+                                      tensor_name):
+  """Overrides initialization op of given variable or list of variables.
+
+  Calls `_set_checkpoint_initializer` for each variable in the given list of
+  variables.
+
+  Args:
+    variable_or_list: `tf.Variable` object or a list of `tf.Variable` objects.
+    ckpt_file: string, full path of the checkpoint.
+    tensor_name: Name of the tensor to load from the checkpoint.
+
+  Raises:
+    ValueError: if all objects in `variable_or_list` are not partitions of the
+      same large variable.
+  """
+  if isinstance(variable_or_list, (list, tuple)):
+    # A set of slices.
+    slice_name = None
+    for v in variable_or_list:
+      slice_info = v._save_slice_info  # pylint:disable=protected-access
+      if slice_name is None:
+        slice_name = slice_info.full_name
+      elif slice_name != slice_info.full_name:
+        raise ValueError("Slices must all be from the same tensor: %s != %s" %
+                         (slice_name, slice_info.full_name))
+      _set_checkpoint_initializer(v, ckpt_file, tensor_name, slice_info.spec)
+  else:
+    _set_checkpoint_initializer(variable_or_list, ckpt_file, tensor_name, "")
+
+
+def _is_variable(x):
+  return (isinstance(x, variables.Variable) or
+          resource_variable_ops.is_resource_variable(x))
+
+
+def _collect_partitioned_variable(name, all_vars):
+  """Returns list of `tf.Variable` that comprise the partitioned variable."""
+  if name + "/part_0" in all_vars:
+    var = []
+    i = 0
+    while name + "/part_%d" % i in all_vars:
+      var.append(all_vars[name + "/part_%d" % i])
+      i += 1
+    return var
+  return None

tensorflow/python/training/checkpoint_utils_iterator_test.py

@@ -0,0 +1,125 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Tests for checkpoints tools."""
+
+import os
+import pathlib
+
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.checkpoint import checkpoint as trackable_utils
+from tensorflow.python.client import session as session_lib
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import gfile
+from tensorflow.python.platform import test
+from tensorflow.python.training import checkpoint_utils
+from tensorflow.python.training import saver as saver_lib
+
+
+class CheckpointIteratorTest(test.TestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def testReturnsEmptyIfNoCheckpointsFound(self):
+    checkpoint_dir = os.path.join(self.get_temp_dir(), "no_checkpoints_found")
+
+    num_found = 0
+    for _ in checkpoint_utils.checkpoints_iterator(checkpoint_dir, timeout=0):
+      num_found += 1
+    self.assertEqual(num_found, 0)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testReturnsSingleCheckpointIfOneCheckpointFound(self):
+    checkpoint_dir = os.path.join(self.get_temp_dir(), "one_checkpoint_found")
+    if not gfile.Exists(checkpoint_dir):
+      gfile.MakeDirs(checkpoint_dir)
+
+    save_path = os.path.join(checkpoint_dir, "model.ckpt")
+
+    a = resource_variable_ops.ResourceVariable(5)
+    self.evaluate(a.initializer)
+    checkpoint = trackable_utils.Checkpoint(a=a)
+    checkpoint.save(file_prefix=save_path)
+
+    num_found = 0
+    for _ in checkpoint_utils.checkpoints_iterator(checkpoint_dir, timeout=0):
+      num_found += 1
+    self.assertEqual(num_found, 1)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testWorksWithFSPath(self):
+    checkpoint_dir = pathlib.Path(self.get_temp_dir()) / "one_checkpoint_found"
+    if not gfile.Exists(checkpoint_dir):
+      gfile.MakeDirs(checkpoint_dir)
+
+    save_path = checkpoint_dir / "model.ckpt"
+
+    a = resource_variable_ops.ResourceVariable(5)
+    self.evaluate(a.initializer)
+    checkpoint = trackable_utils.Checkpoint(a=a)
+    checkpoint.save(file_prefix=save_path)
+
+    num_found = 0
+    for _ in checkpoint_utils.checkpoints_iterator(checkpoint_dir, timeout=0):
+      num_found += 1
+    self.assertEqual(num_found, 1)
+
+  @test_util.run_v1_only("Tests v1-style checkpoint sharding")
+  def testReturnsSingleCheckpointIfOneShardedCheckpoint(self):
+    checkpoint_dir = os.path.join(self.get_temp_dir(),
+                                  "one_checkpoint_found_sharded")
+    if not gfile.Exists(checkpoint_dir):
+      gfile.MakeDirs(checkpoint_dir)
+
+    global_step = variables.Variable(0, name="v0")
+
+    # This will result in 3 different checkpoint shard files.
+    with ops.device("/cpu:0"):
+      variables.Variable(10, name="v1")
+    with ops.device("/cpu:1"):
+      variables.Variable(20, name="v2")
+
+    saver = saver_lib.Saver(sharded=True)
+
+    with session_lib.Session(
+        target="",
+        config=config_pb2.ConfigProto(device_count={"CPU": 2})) as session:
+
+      session.run(variables.global_variables_initializer())
+      save_path = os.path.join(checkpoint_dir, "model.ckpt")
+      saver.save(session, save_path, global_step=global_step)
+
+    num_found = 0
+    for _ in checkpoint_utils.checkpoints_iterator(checkpoint_dir, timeout=0):
+      num_found += 1
+    self.assertEqual(num_found, 1)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testTimeoutFn(self):
+    timeout_fn_calls = [0]
+    def timeout_fn():
+      timeout_fn_calls[0] += 1
+      return timeout_fn_calls[0] > 3
+
+    results = list(
+        checkpoint_utils.checkpoints_iterator(
+            "/non-existent-dir", timeout=0.1, timeout_fn=timeout_fn))
+    self.assertEqual([], results)
+    self.assertEqual(4, timeout_fn_calls[0])
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/checkpoint_utils_test.py

@@ -0,0 +1,430 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Tests for checkpoints tools."""
+
+import os
+import pathlib
+import time
+
+import numpy as np
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import partitioned_variables
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import checkpoint_utils
+from tensorflow.python.training import saver as saver_lib
+
+
+def _create_checkpoints(sess, checkpoint_dir):
+  checkpoint_prefix = os.path.join(checkpoint_dir, "model")
+  checkpoint_state_name = "checkpoint"
+  v1 = variable_scope.get_variable("var1", [1, 10])
+  v2 = variable_scope.get_variable("var2", [10, 10])
+  v3 = variable_scope.get_variable("var3", [100, 100])
+  with variable_scope.variable_scope("useful_scope"):
+    v4 = variable_scope.get_variable("var4", [9, 9])
+  sess.run(variables.global_variables_initializer())
+  v1_value, v2_value, v3_value, v4_value = sess.run([v1, v2, v3, v4])
+  saver = saver_lib.Saver()
+  saver.save(
+      sess,
+      checkpoint_prefix,
+      global_step=0,
+      latest_filename=checkpoint_state_name)
+  return v1_value, v2_value, v3_value, v4_value
+
+
+def _create_partition_checkpoints(sess, checkpoint_dir):
+  checkpoint_prefix = os.path.join(checkpoint_dir, "model")
+  checkpoint_state_name = "checkpoint"
+  with variable_scope.variable_scope("scope"):
+    v1 = variable_scope.get_variable(
+        name="var1",
+        shape=[100, 100],
+        initializer=init_ops.truncated_normal_initializer(0.5),
+        partitioner=partitioned_variables.min_max_variable_partitioner(
+            max_partitions=5, axis=0, min_slice_size=8 << 10))
+  sess.run(variables.global_variables_initializer())
+  v1_value = sess.run(v1._get_variable_list())
+  saver = saver_lib.Saver()
+  saver.save(
+      sess,
+      checkpoint_prefix,
+      global_step=0,
+      latest_filename=checkpoint_state_name)
+  return v1_value
+
+
+class CheckpointsTest(test.TestCase):
+
+  def testNoCheckpoints(self):
+    checkpoint_dir = self.get_temp_dir() + "/no_checkpoints"
+    with self.assertRaises(errors_impl.OpError):
+      self.assertAllEqual(
+          checkpoint_utils.load_variable(checkpoint_dir, "var1"), [])
+
+  def testNoTensor(self):
+    checkpoint_dir = self.get_temp_dir()
+    with self.cached_session() as session:
+      _, _, _, _ = _create_checkpoints(session, checkpoint_dir)
+    with self.assertRaises(errors_impl.OpError):
+      self.assertAllEqual(
+          checkpoint_utils.load_variable(checkpoint_dir, "var5"), [])
+
+  def testGetTensor(self):
+    checkpoint_dir = self.get_temp_dir()
+    with self.cached_session() as session:
+      v1, v2, v3, v4 = _create_checkpoints(session, checkpoint_dir)
+    self.assertAllEqual(
+        checkpoint_utils.load_variable(checkpoint_dir, "var1"), v1)
+    self.assertAllEqual(
+        checkpoint_utils.load_variable(checkpoint_dir, "var2"), v2)
+    self.assertAllEqual(
+        checkpoint_utils.load_variable(checkpoint_dir, "var3"), v3)
+    self.assertAllEqual(
+        checkpoint_utils.load_variable(checkpoint_dir, "useful_scope/var4"), v4)
+
+  def testGetAllVariables(self):
+    checkpoint_dir = self.get_temp_dir()
+    with self.cached_session() as session:
+      _create_checkpoints(session, checkpoint_dir)
+    self.assertEqual(
+        checkpoint_utils.list_variables(checkpoint_dir),
+        [("useful_scope/var4", [9, 9]), ("var1", [1, 10]), ("var2", [10, 10]),
+         ("var3", [100, 100])])
+
+  def testFSPath(self):
+    checkpoint_dir = pathlib.Path(self.get_temp_dir())
+    with self.cached_session() as session:
+      v1, v2, v3, v4 = _create_checkpoints(session, checkpoint_dir)  # pylint: disable=unused-variable
+
+    reader = checkpoint_utils.load_checkpoint(checkpoint_dir)
+    self.assertAllEqual(reader.get_tensor("var1"), v1)
+
+    self.assertAllEqual(
+        checkpoint_utils.load_variable(checkpoint_dir, "var1"), v1)
+
+    self.assertEqual(
+        checkpoint_utils.list_variables(checkpoint_dir),
+        [("useful_scope/var4", [9, 9]), ("var1", [1, 10]), ("var2", [10, 10]),
+         ("var3", [100, 100])])
+
+  def testInitFromCheckpoint(self):
+    checkpoint_dir = self.get_temp_dir()
+    with self.cached_session() as session:
+      v1, v2, v3, v4 = _create_checkpoints(session, checkpoint_dir)
+
+    # New graph and session.
+    with ops.Graph().as_default() as g:
+      with self.session(graph=g) as session:
+        with variable_scope.variable_scope("some_scope"):
+          my1 = variable_scope.get_variable("my1", [1, 10])
+          with variable_scope.variable_scope("some_other_scope"):
+            my2 = variable_scope.get_variable("my2", [10, 10])
+            with variable_scope.variable_scope("other_useful_scope"):
+              my4 = variable_scope.get_variable("var4", [9, 9])
+        my3 = variable_scope.get_variable("my3", [100, 100])
+        my3b = variable_scope.get_variable("my3b", [100, 100])
+
+        checkpoint_utils.init_from_checkpoint(checkpoint_dir, {
+            "var1": "some_scope/my1",
+            "useful_scope/": "some_scope/some_other_scope/other_useful_scope/",
+        })
+        checkpoint_utils.init_from_checkpoint(checkpoint_dir, [
+            ("var2", "some_scope/some_other_scope/my2"),
+            ("var3", my3),
+            ("var3", my3b),
+        ])
+
+        session.run(variables.global_variables_initializer())
+        self.assertAllEqual(my1.eval(session), v1)
+        self.assertAllEqual(my2.eval(session), v2)
+        self.assertAllEqual(my3.eval(session), v3)
+        self.assertAllEqual(my3b.eval(session), v3)
+        self.assertAllEqual(my4.eval(session), v4)
+
+        # Check that tensors are not explicitly in the graph.
+        self.assertLess(len(str(session.graph.as_graph_def())), 32000)
+
+  def testInitialValueComesFromCheckpoint(self):
+    checkpoint_dir = self.get_temp_dir()
+    with self.cached_session() as session:
+      v1, _, _, _ = _create_checkpoints(session, checkpoint_dir)
+
+    # New graph and session.
+    with ops.Graph().as_default() as g:
+      with self.session(graph=g) as session:
+        with variable_scope.variable_scope(
+            "some_scope", initializer=init_ops.zeros_initializer()):
+          my1 = variable_scope.get_variable("my1", [1, 10])
+
+        before = my1.initialized_value()
+
+        checkpoint_utils.init_from_checkpoint(checkpoint_dir, {"var1": my1})
+
+        after = my1.initialized_value()
+
+        self.assertAllEqual(session.run(before), [[0.0] * 10])
+        self.assertAllEqual(session.run(after), v1)
+
+        session.run(variables.global_variables_initializer())
+
+        self.assertAllEqual(session.run(my1), v1)
+        self.assertAllEqual(session.run(my1.initialized_value()), v1)
+        self.assertAllClose(session.run(before), v1)
+        self.assertAllClose(session.run(after), v1)
+        with self.assertRaises(AssertionError):
+          self.assertAllClose(v1, [[0.0] * 10])
+
+  def testInitWithScopeDoesNotCaptureSuffixes(self):
+    checkpoint_dir = self.get_temp_dir()
+    with self.cached_session() as session:
+      _, _, _, v4 = _create_checkpoints(session, checkpoint_dir)
+
+    with ops.Graph().as_default() as g:
+      with variable_scope.variable_scope("useful_scope"):
+        my4 = variable_scope.get_variable("var4", [9, 9])
+      with variable_scope.variable_scope("useful_scope_1"):
+        my5_init = [[1.0, 2.0], [3.0, 4.0]]
+        my5 = variable_scope.get_variable("var5", initializer=my5_init)
+
+      checkpoint_utils.init_from_checkpoint(checkpoint_dir,
+                                            {"useful_scope/": "useful_scope/"})
+      with self.session(graph=g) as session:
+        session.run(variables.global_variables_initializer())
+        self.assertAllEqual(my4.eval(session), v4)
+        self.assertAllEqual(my5.eval(session), my5_init)
+
+  def testRestoreRunsOnSameDevice(self):
+    checkpoint_dir = self.get_temp_dir()
+    with self.cached_session() as session:
+      _create_checkpoints(session, checkpoint_dir)
+
+    with ops.Graph().as_default():
+      with ops.device("/job:ps"):
+        with variable_scope.variable_scope("useful_scope"):
+          variable_scope.get_variable("var4", [9, 9])
+
+      checkpoint_utils.init_from_checkpoint(checkpoint_dir,
+                                            {"useful_scope/": "useful_scope/"})
+
+  def testInitFromRootCheckpoint(self):
+    checkpoint_dir = self.get_temp_dir()
+    with self.cached_session() as session:
+      v1, v2, v3, v4 = _create_checkpoints(session, checkpoint_dir)
+
+    # New graph and session.
+    with ops.Graph().as_default() as g:
+      with self.session(graph=g) as session:
+        with variable_scope.variable_scope("some_scope"):
+          my1 = variable_scope.get_variable("var1", [1, 10])
+          my2 = variable_scope.get_variable("var2", [10, 10])
+          my3 = variable_scope.get_variable("var3", [100, 100])
+          with variable_scope.variable_scope("useful_scope"):
+            my4 = variable_scope.get_variable("var4", [9, 9])
+
+        checkpoint_utils.init_from_checkpoint(checkpoint_dir,
+                                              {"/": "some_scope/",})
+
+        session.run(variables.global_variables_initializer())
+        self.assertAllEqual(my1.eval(session), v1)
+        self.assertAllEqual(my2.eval(session), v2)
+        self.assertAllEqual(my3.eval(session), v3)
+        self.assertAllEqual(my4.eval(session), v4)
+
+  def testInitToRootCheckpoint(self):
+    checkpoint_dir = self.get_temp_dir()
+    with self.cached_session() as session:
+      v1, v2, v3, v4 = _create_checkpoints(session, checkpoint_dir)
+
+    # New graph and session.
+    with ops.Graph().as_default() as g:
+      with self.session(graph=g) as session:
+        my1 = variable_scope.get_variable("var1", [1, 10])
+        my2 = variable_scope.get_variable("var2", [10, 10])
+        my3 = variable_scope.get_variable("var3", [100, 100])
+        with variable_scope.variable_scope("useful_scope"):
+          my4 = variable_scope.get_variable("var4", [9, 9])
+
+        checkpoint_utils.init_from_checkpoint(checkpoint_dir,
+                                              {"/": "/",})
+
+        session.run(variables.global_variables_initializer())
+        self.assertAllEqual(my1.eval(session), v1)
+        self.assertAllEqual(my2.eval(session), v2)
+        self.assertAllEqual(my3.eval(session), v3)
+        self.assertAllEqual(my4.eval(session), v4)
+
+  def testInitFromPartitionVar(self):
+    checkpoint_dir = self.get_temp_dir()
+    with self.cached_session() as session:
+      v1 = _create_partition_checkpoints(session, checkpoint_dir)
+
+    # New graph and session.
+    with ops.Graph().as_default() as g:
+      with self.session(graph=g) as session:
+        with variable_scope.variable_scope("some_scope"):
+          my1 = variable_scope.get_variable(
+              name="my1",
+              shape=[100, 100],
+              initializer=init_ops.zeros_initializer(),
+              partitioner=partitioned_variables.min_max_variable_partitioner(
+                  max_partitions=5, axis=0, min_slice_size=8 << 10))
+          my1_var_list = my1._get_variable_list()
+        # Create another variable with different partitions than the variable in
+        # the checkpoint.
+        with variable_scope.variable_scope("some_other_scope"):
+          my2 = variable_scope.get_variable(
+              name="var1",
+              shape=[100, 100],
+              initializer=init_ops.zeros_initializer(),
+              partitioner=partitioned_variables.min_max_variable_partitioner(
+                  max_partitions=5, axis=0, min_slice_size=16 << 10))
+          my2_var_list = my2._get_variable_list()
+
+        checkpoint_utils.init_from_checkpoint(checkpoint_dir, {
+            "scope/var1": "some_scope/my1",
+            "scope/": "some_other_scope/"})
+
+        session.run(variables.global_variables_initializer())
+        my1_values = session.run(my1_var_list)
+        self.assertAllEqual(my1_values, v1)
+        my2_values = session.run(my2_var_list)
+        # Verify we created different number of partitions.
+        self.assertNotEqual(len(my2_values), len(v1))
+        # Verify the values were correctly initialized inspite of different
+        # partitions.
+        full_my2_values = np.concatenate(my2_values, axis=0)
+        full_v1_values = np.concatenate(v1, axis=0)
+        self.assertAllEqual(full_my2_values, full_v1_values)
+
+    # New graph and session.
+    with ops.Graph().as_default() as g:
+      with self.session(graph=g) as session:
+        with variable_scope.variable_scope("some_scope"):
+          my1 = variable_scope.get_variable(
+              name="my1",
+              shape=[100, 100],
+              initializer=init_ops.truncated_normal_initializer(0.5),
+              partitioner=partitioned_variables.min_max_variable_partitioner(
+                  max_partitions=5, axis=0, min_slice_size=8 << 10))
+          my1_var_list = my1._get_variable_list()
+
+        checkpoint_utils.init_from_checkpoint(checkpoint_dir,
+                                              {"scope/var1": my1_var_list,})
+
+        session.run(variables.global_variables_initializer())
+        my1_values = session.run(my1_var_list)
+        self.assertAllEqual(my1_values, v1)
+
+  def testInitFromCheckpointMissing(self):
+    checkpoint_dir = self.get_temp_dir()
+    with self.cached_session() as session:
+      _, _, _, _ = _create_checkpoints(session, checkpoint_dir)
+
+    # New graph and session.
+    with ops.Graph().as_default() as g:
+      with self.session(graph=g) as session:
+        with variable_scope.variable_scope("some_scope"):
+          _ = variable_scope.get_variable("my1", [10, 10])
+          _ = variable_scope.get_variable(
+              "my2", [1, 10],
+              dtype=dtypes.int64,
+              initializer=init_ops.zeros_initializer())
+
+        # No directory.
+        with self.assertRaises(errors_impl.OpError):
+          checkpoint_utils.init_from_checkpoint("no_dir",
+                                                {"var1": "some_scope/my1"})
+
+        # No variable in checkpoint.
+        with self.assertRaises(ValueError):
+          checkpoint_utils.init_from_checkpoint(checkpoint_dir,
+                                                {"no_var": "some_scope/my1"})
+
+        # No variable in the graph.
+        with self.assertRaises(ValueError):
+          checkpoint_utils.init_from_checkpoint(checkpoint_dir,
+                                                {"var3": "some_scope/no_var"})
+
+        # Shape mismatch.
+        with self.assertRaises(ValueError):
+          checkpoint_utils.init_from_checkpoint(checkpoint_dir,
+                                                {"var1": "some_scope/my1"})
+
+        # Variable 'my1' and 'my2' are missing in given checkpoint scope.
+        with self.assertRaises(ValueError):
+          checkpoint_utils.init_from_checkpoint(
+              checkpoint_dir, {"useful_scope/": "some_scope/"})
+
+        # Mapping is not to scope name.
+        with self.assertRaises(ValueError):
+          checkpoint_utils.init_from_checkpoint(checkpoint_dir,
+                                                {"useful_scope": "some_scope/"})
+
+  def testNoAdditionalReadOpsForResourceVariables(self):
+    checkpoint_dir = self.get_temp_dir()
+    with self.cached_session() as session:
+      v1, _, _, _ = _create_checkpoints(session, checkpoint_dir)
+
+    # New graph and session.
+    with ops.Graph().as_default() as g:
+      with self.session(graph=g) as session:
+        my1 = resource_variable_ops.ResourceVariable([[0.0] * 10], name="my1")
+
+        with ops.name_scope("init_from_checkpoint"):
+          checkpoint_utils.init_from_checkpoint(checkpoint_dir, {"var1": my1})
+
+        # Basic sanity checks:
+        session.run(variables.global_variables_initializer())
+        self.assertAllEqual(session.run(my1), v1)
+
+    ops_in_init_from_checkpoint_scope = [
+        op for op in g.get_operations()
+        if (op.name.startswith("init_from_checkpoint/") and
+            not op.name.startswith("init_from_checkpoint/checkpoint_initializer"
+                                  ) and
+            op.type != "AssignVariableOp" and
+            op.type != "Identity")
+    ]
+    self.assertEqual(ops_in_init_from_checkpoint_scope, [])
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class WaitForNewCheckpointTest(test.TestCase):
+
+  def testReturnsNoneAfterTimeout(self):
+    start = time.time()
+    ret = checkpoint_utils.wait_for_new_checkpoint(
+        "/non-existent-dir", "foo", timeout=1.0, seconds_to_sleep=0.5)
+    end = time.time()
+    self.assertIsNone(ret)
+
+    # We've waited one second.
+    self.assertGreater(end, start + 0.5)
+
+    # The timeout kicked in.
+    self.assertLess(end, start + 1.1)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/coordinator.py

@@ -0,0 +1,507 @@
+# 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.
+# ==============================================================================
+"""Coordinator to help multiple threads stop when requested."""
+import contextlib
+import sys
+import threading
+import time
+
+from tensorflow.python.framework import errors
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.util import compat
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("train.Coordinator")
+class Coordinator:
+  """A coordinator for threads.
+
+  This class implements a simple mechanism to coordinate the termination of a
+  set of threads.
+
+  #### Usage:
+
+  ```python
+  # Create a coordinator.
+  coord = Coordinator()
+  # Start a number of threads, passing the coordinator to each of them.
+  ...start thread 1...(coord, ...)
+  ...start thread N...(coord, ...)
+  # Wait for all the threads to terminate.
+  coord.join(threads)
+  ```
+
+  Any of the threads can call `coord.request_stop()` to ask for all the threads
+  to stop.  To cooperate with the requests, each thread must check for
+  `coord.should_stop()` on a regular basis.  `coord.should_stop()` returns
+  `True` as soon as `coord.request_stop()` has been called.
+
+  A typical thread running with a coordinator will do something like:
+
+  ```python
+  while not coord.should_stop():
+    ...do some work...
+  ```
+
+  #### Exception handling:
+
+  A thread can report an exception to the coordinator as part of the
+  `request_stop()` call.  The exception will be re-raised from the
+  `coord.join()` call.
+
+  Thread code:
+
+  ```python
+  try:
+    while not coord.should_stop():
+      ...do some work...
+  except Exception as e:
+    coord.request_stop(e)
+  ```
+
+  Main code:
+
+  ```python
+  try:
+    ...
+    coord = Coordinator()
+    # Start a number of threads, passing the coordinator to each of them.
+    ...start thread 1...(coord, ...)
+    ...start thread N...(coord, ...)
+    # Wait for all the threads to terminate.
+    coord.join(threads)
+  except Exception as e:
+    ...exception that was passed to coord.request_stop()
+  ```
+
+  To simplify the thread implementation, the Coordinator provides a
+  context handler `stop_on_exception()` that automatically requests a stop if
+  an exception is raised.  Using the context handler the thread code above
+  can be written as:
+
+  ```python
+  with coord.stop_on_exception():
+    while not coord.should_stop():
+      ...do some work...
+  ```
+
+  #### Grace period for stopping:
+
+  After a thread has called `coord.request_stop()` the other threads have a
+  fixed time to stop, this is called the 'stop grace period' and defaults to 2
+  minutes.  If any of the threads is still alive after the grace period expires
+  `coord.join()` raises a RuntimeError reporting the laggards.
+
+  ```python
+  try:
+    ...
+    coord = Coordinator()
+    # Start a number of threads, passing the coordinator to each of them.
+    ...start thread 1...(coord, ...)
+    ...start thread N...(coord, ...)
+    # Wait for all the threads to terminate, give them 10s grace period
+    coord.join(threads, stop_grace_period_secs=10)
+  except RuntimeError:
+    ...one of the threads took more than 10s to stop after request_stop()
+    ...was called.
+  except Exception:
+    ...exception that was passed to coord.request_stop()
+  ```
+  """
+
+  def __init__(self, clean_stop_exception_types=None):
+    """Create a new Coordinator.
+
+    Args:
+      clean_stop_exception_types: Optional tuple of Exception types that should
+        cause a clean stop of the coordinator. If an exception of one of these
+        types is reported to `request_stop(ex)` the coordinator will behave as
+        if `request_stop(None)` was called.  Defaults to
+        `(tf.errors.OutOfRangeError,)` which is used by input queues to signal
+        the end of input. When feeding training data from a Python iterator it
+        is common to add `StopIteration` to this list.
+    """
+    if clean_stop_exception_types is None:
+      clean_stop_exception_types = (errors.OutOfRangeError,)
+    self._clean_stop_exception_types = tuple(clean_stop_exception_types)
+    # Protects all attributes.
+    self._lock = threading.Lock()
+    # Event set when threads must stop.
+    self._stop_event = threading.Event()
+    # Python exc_info to report.
+    # If not None, it should hold the returned value of sys.exc_info(), which is
+    # a tuple containing exception (type, value, traceback).
+    self._exc_info_to_raise = None
+    # True if we have called join() already.
+    self._joined = False
+    # Set of threads registered for joining when join() is called.  These
+    # threads will be joined in addition to the threads passed to the join()
+    # call.  It's ok if threads are both registered and passed to the join()
+    # call.
+    self._registered_threads = set()
+
+  def _filter_exception(self, ex):
+    """Check if the exception indicated in 'ex' should be ignored.
+
+    This method examines `ex` to check if it is an exception that should be
+    reported to the users.  If yes, it returns `ex` as is, otherwise it returns
+    None.
+
+    The code returns None for exception types listed in
+    `_clean_stop_exception_types`.
+
+    Args:
+      ex: None, an `Exception`, or a Python `exc_info` tuple as returned by
+        `sys.exc_info()`.
+
+    Returns:
+      ex or None.
+    """
+    if isinstance(ex, tuple):
+      ex2 = ex[1]
+    else:
+      ex2 = ex
+    if isinstance(ex2, self._clean_stop_exception_types):
+      # Ignore the exception.
+      ex = None
+    return ex
+
+  def request_stop(self, ex=None):
+    """Request that the threads stop.
+
+    After this is called, calls to `should_stop()` will return `True`.
+
+    Note: If an exception is being passed in, in must be in the context of
+    handling the exception (i.e. `try: ... except Exception as ex: ...`) and not
+    a newly created one.
+
+    Args:
+      ex: Optional `Exception`, or Python `exc_info` tuple as returned by
+        `sys.exc_info()`.  If this is the first call to `request_stop()` the
+        corresponding exception is recorded and re-raised from `join()`.
+    """
+    with self._lock:
+      ex = self._filter_exception(ex)
+      # If we have already joined the coordinator the exception will not have a
+      # chance to be reported, so just raise it normally.  This can happen if
+      # you continue to use a session have having stopped and joined the
+      # coordinator threads.
+      if self._joined:
+        if isinstance(ex, tuple):
+          _, ex_instance, _ = ex
+          raise ex_instance
+        elif ex is not None:
+          # NOTE(touts): This is bogus if request_stop() is not called
+          # from the exception handler that raised ex.
+          _, ex_instance, _ = sys.exc_info()
+          raise ex_instance
+      if not self._stop_event.is_set():
+        if ex and self._exc_info_to_raise is None:
+          if isinstance(ex, tuple):
+            logging.info("Error reported to Coordinator: %s",
+                         compat.as_str_any(ex[1]),
+                         exc_info=ex)
+            self._exc_info_to_raise = ex
+          else:
+            logging.info("Error reported to Coordinator: %s, %s",
+                         type(ex),
+                         compat.as_str_any(ex))
+            self._exc_info_to_raise = sys.exc_info()
+          # self._exc_info_to_raise should contain a tuple containing exception
+          # (type, value, traceback)
+          if (len(self._exc_info_to_raise) != 3 or
+              not self._exc_info_to_raise[0] or
+              not self._exc_info_to_raise[1]):
+            # Raise, catch and record the exception here so that error happens
+            # where expected.
+            try:
+              raise ValueError(
+                  "ex must be a tuple or sys.exc_info must return the current "
+                  "exception: %s"
+                  % self._exc_info_to_raise)
+            except ValueError:
+              # Record this error so it kills the coordinator properly.
+              # NOTE(touts): As above, this is bogus if request_stop() is not
+              # called from the exception handler that raised ex.
+              self._exc_info_to_raise = sys.exc_info()
+
+        self._stop_event.set()
+
+  def clear_stop(self):
+    """Clears the stop flag.
+
+    After this is called, calls to `should_stop()` will return `False`.
+    """
+    with self._lock:
+      self._joined = False
+      self._exc_info_to_raise = None
+      if self._stop_event.is_set():
+        self._stop_event.clear()
+
+  def should_stop(self):
+    """Check if stop was requested.
+
+    Returns:
+      True if a stop was requested.
+    """
+    return self._stop_event.is_set()
+
+  @contextlib.contextmanager
+  def stop_on_exception(self):
+    """Context manager to request stop when an Exception is raised.
+
+    Code that uses a coordinator must catch exceptions and pass
+    them to the `request_stop()` method to stop the other threads
+    managed by the coordinator.
+
+    This context handler simplifies the exception handling.
+    Use it as follows:
+
+    ```python
+    with coord.stop_on_exception():
+      # Any exception raised in the body of the with
+      # clause is reported to the coordinator before terminating
+      # the execution of the body.
+      ...body...
+    ```
+
+    This is completely equivalent to the slightly longer code:
+
+    ```python
+    try:
+      ...body...
+    except:
+      coord.request_stop(sys.exc_info())
+    ```
+
+    Yields:
+      nothing.
+    """
+    try:
+      yield
+    except:  # pylint: disable=bare-except
+      self.request_stop(ex=sys.exc_info())
+
+  def wait_for_stop(self, timeout=None):
+    """Wait till the Coordinator is told to stop.
+
+    Args:
+      timeout: Float.  Sleep for up to that many seconds waiting for
+        should_stop() to become True.
+
+    Returns:
+      True if the Coordinator is told stop, False if the timeout expired.
+    """
+    return self._stop_event.wait(timeout)
+
+  def register_thread(self, thread):
+    """Register a thread to join.
+
+    Args:
+      thread: A Python thread to join.
+    """
+    with self._lock:
+      self._registered_threads.add(thread)
+
+  def join(self, threads=None, stop_grace_period_secs=120,
+           ignore_live_threads=False):
+    """Wait for threads to terminate.
+
+    This call blocks until a set of threads have terminated.  The set of thread
+    is the union of the threads passed in the `threads` argument and the list
+    of threads that registered with the coordinator by calling
+    `Coordinator.register_thread()`.
+
+    After the threads stop, if an `exc_info` was passed to `request_stop`, that
+    exception is re-raised.
+
+    Grace period handling: When `request_stop()` is called, threads are given
+    'stop_grace_period_secs' seconds to terminate.  If any of them is still
+    alive after that period expires, a `RuntimeError` is raised.  Note that if
+    an `exc_info` was passed to `request_stop()` then it is raised instead of
+    that `RuntimeError`.
+
+    Args:
+      threads: List of `threading.Threads`. The started threads to join in
+        addition to the registered threads.
+      stop_grace_period_secs: Number of seconds given to threads to stop after
+        `request_stop()` has been called.
+      ignore_live_threads: If `False`, raises an error if any of the threads are
+        still alive after `stop_grace_period_secs`.
+
+    Raises:
+      RuntimeError: If any thread is still alive after `request_stop()`
+        is called and the grace period expires.
+    """
+    # Threads registered after this call will not be joined.
+    with self._lock:
+      if threads is None:
+        threads = self._registered_threads
+      else:
+        threads = self._registered_threads.union(set(threads))
+      # Copy the set into a list to avoid race conditions where a new thread
+      # is added while we are waiting.
+      threads = list(threads)
+
+    # Wait for all threads to stop or for request_stop() to be called.
+    while any(t.is_alive() for t in threads) and not self.wait_for_stop(1.0):
+      pass
+
+    # If any thread is still alive, wait for the grace period to expire.
+    # By the time this check is executed, threads may still be shutting down,
+    # so we add a sleep of increasing duration to give them a chance to shut
+    # down without losing too many cycles.
+    # The sleep duration is limited to the remaining grace duration.
+    stop_wait_secs = 0.001
+    while any(t.is_alive() for t in threads) and stop_grace_period_secs >= 0.0:
+      time.sleep(stop_wait_secs)
+      stop_grace_period_secs -= stop_wait_secs
+      stop_wait_secs = 2 * stop_wait_secs
+      # Keep the waiting period within sane bounds.
+      # The minimum value is to avoid decreasing stop_wait_secs to a value
+      # that could cause stop_grace_period_secs to remain unchanged.
+      stop_wait_secs = max(min(stop_wait_secs, stop_grace_period_secs), 0.001)
+
+    # List the threads still alive after the grace period.
+    stragglers = [t.name for t in threads if t.is_alive()]
+
+    # Terminate with an exception if appropriate.
+    with self._lock:
+      self._joined = True
+      self._registered_threads = set()
+      if self._exc_info_to_raise:
+        _, ex_instance, _ = self._exc_info_to_raise
+        raise ex_instance
+      elif stragglers:
+        if ignore_live_threads:
+          logging.info("Coordinator stopped with threads still running: %s",
+                       " ".join(stragglers))
+        else:
+          raise RuntimeError(
+              "Coordinator stopped with threads still running: %s" %
+              " ".join(stragglers))
+
+  @property
+  def joined(self):
+    return self._joined
+
+  def raise_requested_exception(self):
+    """If an exception has been passed to `request_stop`, this raises it."""
+    with self._lock:
+      if self._exc_info_to_raise:
+        _, ex_instance, _ = self._exc_info_to_raise
+        raise ex_instance
+
+
+# Threads for the standard services.
+@tf_export(v1=["train.LooperThread"])
+class LooperThread(threading.Thread):
+  """A thread that runs code repeatedly, optionally on a timer.
+
+  This thread class is intended to be used with a `Coordinator`.  It repeatedly
+  runs code specified either as `target` and `args` or by the `run_loop()`
+  method.
+
+  Before each run the thread checks if the coordinator has requested stop.  In
+  that case the looper thread terminates immediately.
+
+  If the code being run raises an exception, that exception is reported to the
+  coordinator and the thread terminates.  The coordinator will then request all
+  the other threads it coordinates to stop.
+
+  You typically pass looper threads to the supervisor `Join()` method.
+  """
+
+  def __init__(self, coord, timer_interval_secs, target=None, args=None,
+               kwargs=None):
+    """Create a LooperThread.
+
+    Args:
+      coord: A Coordinator.
+      timer_interval_secs: Time boundaries at which to call Run(), or None
+        if it should be called back to back.
+      target: Optional callable object that will be executed in the thread.
+      args: Optional arguments to pass to `target` when calling it.
+      kwargs: Optional keyword arguments to pass to `target` when calling it.
+
+    Raises:
+      ValueError: If one of the arguments is invalid.
+    """
+    if not isinstance(coord, Coordinator):
+      raise ValueError("'coord' argument must be a Coordinator: %s" % coord)
+    super(LooperThread, self).__init__()
+    self.daemon = True
+    self._coord = coord
+    self._timer_interval_secs = timer_interval_secs
+    self._target = target
+    if self._target:
+      self._args = args or ()
+      self._kwargs = kwargs or {}
+    elif args or kwargs:
+      raise ValueError("'args' and 'kwargs' argument require that you also "
+                       "pass 'target'")
+    self._coord.register_thread(self)
+
+  @staticmethod
+  def loop(coord, timer_interval_secs, target, args=None, kwargs=None):
+    """Start a LooperThread that calls a function periodically.
+
+    If `timer_interval_secs` is None the thread calls `target(args)`
+    repeatedly.  Otherwise `target(args)` is called every `timer_interval_secs`
+    seconds.  The thread terminates when a stop of the coordinator is
+    requested.
+
+    Args:
+      coord: A Coordinator.
+      timer_interval_secs: Number. Time boundaries at which to call `target`.
+      target: A callable object.
+      args: Optional arguments to pass to `target` when calling it.
+      kwargs: Optional keyword arguments to pass to `target` when calling it.
+
+    Returns:
+      The started thread.
+    """
+    looper = LooperThread(coord, timer_interval_secs, target=target, args=args,
+                          kwargs=kwargs)
+    looper.start()
+    return looper
+
+  def run(self):
+    with self._coord.stop_on_exception():
+      self.start_loop()
+      if self._timer_interval_secs is None:
+        # Call back-to-back.
+        while not self._coord.should_stop():
+          self.run_loop()
+      else:
+        # Next time at which to call run_loop(), starts as 'now'.
+        next_timer_time = time.time()
+        while not self._coord.wait_for_stop(next_timer_time - time.time()):
+          next_timer_time += self._timer_interval_secs
+          self.run_loop()
+      self.stop_loop()
+
+  def start_loop(self):
+    """Called when the thread starts."""
+    pass
+
+  def stop_loop(self):
+    """Called when the thread stops."""
+    pass
+
+  def run_loop(self):
+    """Called at 'timer_interval_secs' boundaries."""
+    if self._target:
+      self._target(*self._args, **self._kwargs)

tensorflow/python/training/coordinator_test.py

@@ -0,0 +1,373 @@
+# 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.
+# ==============================================================================
+"""Tests for Coordinator."""
+
+import sys
+import threading
+import time
+
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.platform import test
+from tensorflow.python.training import coordinator
+
+
+def StopOnEvent(coord, wait_for_stop, set_when_stopped):
+  wait_for_stop.wait()
+  coord.request_stop()
+  set_when_stopped.set()
+
+
+def RaiseOnEvent(coord, wait_for_stop, set_when_stopped, ex, report_exception):
+  try:
+    wait_for_stop.wait()
+    raise ex
+  except RuntimeError as e:
+    if report_exception:
+      coord.request_stop(e)
+    else:
+      coord.request_stop(sys.exc_info())
+  finally:
+    if set_when_stopped:
+      set_when_stopped.set()
+
+
+def RaiseOnEventUsingContextHandler(coord, wait_for_stop, set_when_stopped, ex):
+  with coord.stop_on_exception():
+    wait_for_stop.wait()
+    raise ex
+  if set_when_stopped:
+    set_when_stopped.set()
+
+
+def SleepABit(n_secs, coord=None):
+  if coord:
+    coord.register_thread(threading.current_thread())
+  time.sleep(n_secs)
+
+
+def WaitForThreadsToRegister(coord, num_threads):
+  while True:
+    with coord._lock:
+      if len(coord._registered_threads) == num_threads:
+        break
+    time.sleep(0.001)
+
+
+class CoordinatorTest(test.TestCase):
+
+  def testStopAPI(self):
+    coord = coordinator.Coordinator()
+    self.assertFalse(coord.should_stop())
+    self.assertFalse(coord.wait_for_stop(0.01))
+    coord.request_stop()
+    self.assertTrue(coord.should_stop())
+    self.assertTrue(coord.wait_for_stop(0.01))
+
+  def testStopAsync(self):
+    coord = coordinator.Coordinator()
+    self.assertFalse(coord.should_stop())
+    self.assertFalse(coord.wait_for_stop(0.1))
+    wait_for_stop_ev = threading.Event()
+    has_stopped_ev = threading.Event()
+    t = threading.Thread(
+        target=StopOnEvent, args=(coord, wait_for_stop_ev, has_stopped_ev))
+    t.start()
+    self.assertFalse(coord.should_stop())
+    self.assertFalse(coord.wait_for_stop(0.01))
+    wait_for_stop_ev.set()
+    has_stopped_ev.wait()
+    self.assertTrue(coord.wait_for_stop(0.05))
+    self.assertTrue(coord.should_stop())
+
+  def testJoin(self):
+    coord = coordinator.Coordinator()
+    threads = [
+        threading.Thread(target=SleepABit, args=(0.01,)),
+        threading.Thread(target=SleepABit, args=(0.02,)),
+        threading.Thread(target=SleepABit, args=(0.01,))
+    ]
+    for t in threads:
+      t.start()
+    coord.join(threads)
+    for t in threads:
+      self.assertFalse(t.is_alive())
+
+  def testJoinAllRegistered(self):
+    coord = coordinator.Coordinator()
+    threads = [
+        threading.Thread(target=SleepABit, args=(0.01, coord)),
+        threading.Thread(target=SleepABit, args=(0.02, coord)),
+        threading.Thread(target=SleepABit, args=(0.01, coord))
+    ]
+    for t in threads:
+      t.start()
+    WaitForThreadsToRegister(coord, 3)
+    coord.join()
+    for t in threads:
+      self.assertFalse(t.is_alive())
+
+  def testJoinSomeRegistered(self):
+    coord = coordinator.Coordinator()
+    threads = [
+        threading.Thread(target=SleepABit, args=(0.01, coord)),
+        threading.Thread(target=SleepABit, args=(0.02,)),
+        threading.Thread(target=SleepABit, args=(0.01, coord))
+    ]
+    for t in threads:
+      t.start()
+    WaitForThreadsToRegister(coord, 2)
+    # threads[1] is not registered we must pass it in.
+    coord.join([threads[1]])
+    for t in threads:
+      self.assertFalse(t.is_alive())
+
+  def testJoinGraceExpires(self):
+
+    def TestWithGracePeriod(stop_grace_period):
+      coord = coordinator.Coordinator()
+      wait_for_stop_ev = threading.Event()
+      has_stopped_ev = threading.Event()
+      threads = [
+          threading.Thread(
+              target=StopOnEvent,
+              args=(coord, wait_for_stop_ev, has_stopped_ev)),
+          threading.Thread(target=SleepABit, args=(10.0,))
+      ]
+      for t in threads:
+        t.daemon = True
+        t.start()
+      wait_for_stop_ev.set()
+      has_stopped_ev.wait()
+      with self.assertRaisesRegex(RuntimeError, "threads still running"):
+        coord.join(threads, stop_grace_period_secs=stop_grace_period)
+
+    TestWithGracePeriod(1e-10)
+    TestWithGracePeriod(0.002)
+    TestWithGracePeriod(1.0)
+
+  def testJoinWithoutGraceExpires(self):
+    coord = coordinator.Coordinator()
+    wait_for_stop_ev = threading.Event()
+    has_stopped_ev = threading.Event()
+    threads = [
+        threading.Thread(
+            target=StopOnEvent, args=(coord, wait_for_stop_ev, has_stopped_ev)),
+        threading.Thread(target=SleepABit, args=(10.0,))
+    ]
+    for t in threads:
+      t.daemon = True
+      t.start()
+    wait_for_stop_ev.set()
+    has_stopped_ev.wait()
+    coord.join(threads, stop_grace_period_secs=1., ignore_live_threads=True)
+
+  def testJoinRaiseReportExcInfo(self):
+    coord = coordinator.Coordinator()
+    ev_1 = threading.Event()
+    ev_2 = threading.Event()
+    threads = [
+        threading.Thread(
+            target=RaiseOnEvent,
+            args=(coord, ev_1, ev_2, RuntimeError("First"), False)),
+        threading.Thread(
+            target=RaiseOnEvent,
+            args=(coord, ev_2, None, RuntimeError("Too late"), False))
+    ]
+    for t in threads:
+      t.start()
+
+    ev_1.set()
+
+    with self.assertRaisesRegex(RuntimeError, "First"):
+      coord.join(threads)
+
+  def testJoinRaiseReportException(self):
+    coord = coordinator.Coordinator()
+    ev_1 = threading.Event()
+    ev_2 = threading.Event()
+    threads = [
+        threading.Thread(
+            target=RaiseOnEvent,
+            args=(coord, ev_1, ev_2, RuntimeError("First"), True)),
+        threading.Thread(
+            target=RaiseOnEvent,
+            args=(coord, ev_2, None, RuntimeError("Too late"), True))
+    ]
+    for t in threads:
+      t.start()
+
+    ev_1.set()
+    with self.assertRaisesRegex(RuntimeError, "First"):
+      coord.join(threads)
+
+  def testJoinIgnoresOutOfRange(self):
+    coord = coordinator.Coordinator()
+    ev_1 = threading.Event()
+    threads = [
+        threading.Thread(
+            target=RaiseOnEvent,
+            args=(coord, ev_1, None,
+                  errors_impl.OutOfRangeError(None, None, "First"), True))
+    ]
+    for t in threads:
+      t.start()
+
+    ev_1.set()
+    coord.join(threads)
+
+  def testJoinIgnoresMyExceptionType(self):
+    coord = coordinator.Coordinator(clean_stop_exception_types=(ValueError,))
+    ev_1 = threading.Event()
+    threads = [
+        threading.Thread(
+            target=RaiseOnEvent,
+            args=(coord, ev_1, None, ValueError("Clean stop"), True))
+    ]
+    for t in threads:
+      t.start()
+
+    ev_1.set()
+    coord.join(threads)
+
+  def testJoinRaiseReportExceptionUsingHandler(self):
+    coord = coordinator.Coordinator()
+    ev_1 = threading.Event()
+    ev_2 = threading.Event()
+    threads = [
+        threading.Thread(
+            target=RaiseOnEventUsingContextHandler,
+            args=(coord, ev_1, ev_2, RuntimeError("First"))),
+        threading.Thread(
+            target=RaiseOnEventUsingContextHandler,
+            args=(coord, ev_2, None, RuntimeError("Too late")))
+    ]
+    for t in threads:
+      t.start()
+
+    ev_1.set()
+    with self.assertRaisesRegex(RuntimeError, "First"):
+      coord.join(threads)
+
+  def testClearStopClearsExceptionToo(self):
+    coord = coordinator.Coordinator()
+    ev_1 = threading.Event()
+    threads = [
+        threading.Thread(
+            target=RaiseOnEvent,
+            args=(coord, ev_1, None, RuntimeError("First"), True)),
+    ]
+    for t in threads:
+      t.start()
+
+    with self.assertRaisesRegex(RuntimeError, "First"):
+      ev_1.set()
+      coord.join(threads)
+    coord.clear_stop()
+    threads = [
+        threading.Thread(
+            target=RaiseOnEvent,
+            args=(coord, ev_1, None, RuntimeError("Second"), True)),
+    ]
+    for t in threads:
+      t.start()
+    with self.assertRaisesRegex(RuntimeError, "Second"):
+      ev_1.set()
+      coord.join(threads)
+
+  def testRequestStopRaisesIfJoined(self):
+    coord = coordinator.Coordinator()
+    # Join the coordinator right away.
+    coord.join([])
+    reported = False
+    with self.assertRaisesRegex(RuntimeError, "Too late"):
+      try:
+        raise RuntimeError("Too late")
+      except RuntimeError as e:
+        reported = True
+        coord.request_stop(e)
+    self.assertTrue(reported)
+    # If we clear_stop the exceptions are handled normally.
+    coord.clear_stop()
+    try:
+      raise RuntimeError("After clear")
+    except RuntimeError as e:
+      coord.request_stop(e)
+    with self.assertRaisesRegex(RuntimeError, "After clear"):
+      coord.join([])
+
+  def testRequestStopRaisesIfJoined_ExcInfo(self):
+    # Same as testRequestStopRaisesIfJoined but using syc.exc_info().
+    coord = coordinator.Coordinator()
+    # Join the coordinator right away.
+    coord.join([])
+    reported = False
+    with self.assertRaisesRegex(RuntimeError, "Too late"):
+      try:
+        raise RuntimeError("Too late")
+      except RuntimeError:
+        reported = True
+        coord.request_stop(sys.exc_info())
+    self.assertTrue(reported)
+    # If we clear_stop the exceptions are handled normally.
+    coord.clear_stop()
+    try:
+      raise RuntimeError("After clear")
+    except RuntimeError:
+      coord.request_stop(sys.exc_info())
+    with self.assertRaisesRegex(RuntimeError, "After clear"):
+      coord.join([])
+
+
+def _StopAt0(coord, n):
+  if n[0] == 0:
+    coord.request_stop()
+  else:
+    n[0] -= 1
+
+
+class LooperTest(test.TestCase):
+
+  def testTargetArgs(self):
+    n = [3]
+    coord = coordinator.Coordinator()
+    thread = coordinator.LooperThread.loop(
+        coord, 0, target=_StopAt0, args=(coord, n))
+    coord.join([thread])
+    self.assertEqual(0, n[0])
+
+  def testTargetKwargs(self):
+    n = [3]
+    coord = coordinator.Coordinator()
+    thread = coordinator.LooperThread.loop(
+        coord, 0, target=_StopAt0, kwargs={
+            "coord": coord,
+            "n": n
+        })
+    coord.join([thread])
+    self.assertEqual(0, n[0])
+
+  def testTargetMixedArgs(self):
+    n = [3]
+    coord = coordinator.Coordinator()
+    thread = coordinator.LooperThread.loop(
+        coord, 0, target=_StopAt0, args=(coord,), kwargs={
+            "n": n
+        })
+    coord.join([thread])
+    self.assertEqual(0, n[0])
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/device_setter.py

@@ -0,0 +1,223 @@
+# 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.
+# ==============================================================================
+"""Device function for replicated training."""
+from tensorflow.core.framework import node_def_pb2
+from tensorflow.python.framework import device as pydev
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.training import server_lib
+from tensorflow.python.util.tf_export import tf_export
+
+STANDARD_PS_OPS = ("Variable", "VariableV2", "AutoReloadVariable",
+                   "MutableHashTable", "MutableHashTableV2",
+                   "MutableHashTableOfTensors", "MutableHashTableOfTensorsV2",
+                   "MutableDenseHashTable", "MutableDenseHashTableV2",
+                   "VarHandleOp", "BoostedTreesEnsembleResourceHandleOp",
+                   "BoostedTreesQuantileStreamResourceHandleOp",
+                   "ResourceConditionalAccumulator",
+                   "DecisionTreeResource")
+
+
+class _RoundRobinStrategy:
+  """Returns the next ps task index for placement in round-robin order.
+
+  This class is not to be used directly by users.  See instead
+  `replica_device_setter()` below.
+  """
+
+  def __init__(self, num_tasks):
+    """Create a new `_RoundRobinStrategy`.
+
+    Args:
+      num_tasks: Number of ps tasks to cycle among.
+    """
+    self._num_tasks = num_tasks
+    self._next_task = 0
+
+  def __call__(self, unused_op):
+    """Choose a ps task index for the given `Operation`.
+
+    Args:
+      unused_op: An `Operation` to be placed on ps.
+
+    Returns:
+      The next ps task index to use for the `Operation`. Returns the next
+      index, in the range `[offset, offset + num_tasks)`.
+    """
+    task = self._next_task
+    self._next_task = (self._next_task + 1) % self._num_tasks
+    return task
+
+
+class _ReplicaDeviceChooser:
+  """Class to choose devices for Ops in a replicated training setup.
+
+  This class is not to be used directly by users.  See instead
+  `replica_device_setter()` below.
+  """
+
+  def __init__(self, ps_tasks, ps_device, worker_device, merge_devices, ps_ops,
+               ps_strategy):
+    """Create a new `_ReplicaDeviceChooser`.
+
+    Args:
+      ps_tasks: Number of tasks in the `ps` job.
+      ps_device: String.  Name of the `ps` job.
+      worker_device: String.  Name of the `worker` job.
+      merge_devices: Boolean. Set to True to allow merging of device specs.
+      ps_ops: List of strings representing `Operation` types that need to be
+        placed on `ps` devices.
+      ps_strategy: A callable invoked for every ps `Operation` (i.e. matched by
+        `ps_ops`), that takes the `Operation` and returns the ps task index to
+        use.
+    """
+    self._ps_tasks = ps_tasks
+    self._ps_device = ps_device
+    self._worker_device = worker_device
+    self._merge_devices = merge_devices
+    self._ps_ops = ps_ops
+    self._ps_strategy = ps_strategy
+
+  def device_function(self, op):
+    """Choose a device for `op`.
+
+    Args:
+      op: an `Operation`.
+
+    Returns:
+      The device to use for the `Operation`.
+    """
+    # If we don't return early here, either merge_devices is True, or op.device
+    # is empty (in which case merging is a no-op). So we can always merge below.
+    if not self._merge_devices and op.device:
+      return op.device
+
+    current_device = pydev.DeviceSpec.from_string(op.device or "")
+
+    # The ps_device will be used for specified ops (ps_ops) whenever it is
+    # present and ps_tasks is non-zero. However, its task number will only be
+    # set (using ps_strategy) if there is a job field in ps_device that won't be
+    # changed by the job field (if present) in current_device.
+    node_def = op if isinstance(op, node_def_pb2.NodeDef) else op.node_def
+    if self._ps_tasks and self._ps_device and node_def.op in self._ps_ops:
+      ps_device = pydev.DeviceSpec.from_string(self._ps_device)
+
+      current_job, ps_job = current_device.job, ps_device.job
+      if ps_job and (not current_job or current_job == ps_job):
+        ps_device = ps_device.replace(task=self._ps_strategy(op))
+
+      ps_device = ps_device.make_merged_spec(current_device)
+      return ps_device.to_string()
+
+    worker_device = pydev.DeviceSpec.from_string(self._worker_device or "")
+    worker_device = worker_device.make_merged_spec(current_device)
+    return worker_device.to_string()
+
+
+@tf_export(v1=["train.replica_device_setter"])
+def replica_device_setter(ps_tasks=0,
+                          ps_device="/job:ps",
+                          worker_device="/job:worker",
+                          merge_devices=True,
+                          cluster=None,
+                          ps_ops=None,
+                          ps_strategy=None):
+  """Return a `device function` to use when building a Graph for replicas.
+
+  Device Functions are used in `with tf.device(device_function):` statement to
+  automatically assign devices to `Operation` objects as they are constructed,
+  Device constraints are added from the inner-most context first, working
+  outwards. The merging behavior adds constraints to fields that are yet unset
+  by a more inner context. Currently the fields are (job, task, cpu/gpu).
+
+  If `cluster` is `None`, and `ps_tasks` is 0, the returned function is a no-op.
+  Otherwise, the value of `ps_tasks` is derived from `cluster`.
+
+  By default, only Variable ops are placed on ps tasks, and the placement
+  strategy is round-robin over all ps tasks. A custom `ps_strategy` may be used
+  to do more intelligent placement, such as
+  `tf.contrib.training.GreedyLoadBalancingStrategy`.
+
+  For example,
+
+  ```python
+  # To build a cluster with two ps jobs on hosts ps0 and ps1, and 3 worker
+  # jobs on hosts worker0, worker1 and worker2.
+  cluster_spec = {
+      "ps": ["ps0:2222", "ps1:2222"],
+      "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]}
+  with
+  tf.compat.v1.device(tf.compat.v1.train.replica_device_setter(cluster=cluster_spec)):
+    # Build your graph
+    v1 = tf.Variable(...)  # assigned to /job:ps/task:0
+    v2 = tf.Variable(...)  # assigned to /job:ps/task:1
+    v3 = tf.Variable(...)  # assigned to /job:ps/task:0
+  # Run compute
+  ```
+
+  Args:
+    ps_tasks: Number of tasks in the `ps` job.  Ignored if `cluster` is
+      provided.
+    ps_device: String.  Device of the `ps` job.  If empty no `ps` job is used.
+      Defaults to `ps`.
+    worker_device: String.  Device of the `worker` job.  If empty no `worker`
+      job is used.
+    merge_devices: `Boolean`. If `True`, merges or only sets a device if the
+      device constraint is completely unset. merges device specification rather
+      than overriding them.
+    cluster: `ClusterDef` proto or `ClusterSpec`.
+    ps_ops: List of strings representing `Operation` types that need to be
+      placed on `ps` devices.  If `None`, defaults to `STANDARD_PS_OPS`.
+    ps_strategy: A callable invoked for every ps `Operation` (i.e. matched by
+      `ps_ops`), that takes the `Operation` and returns the ps task index to
+      use.  If `None`, defaults to a round-robin strategy across all `ps`
+      devices.
+
+  Returns:
+    A function to pass to `tf.device()`.
+
+  Raises:
+    TypeError if `cluster` is not a dictionary or `ClusterDef` protocol buffer,
+    or if `ps_strategy` is provided but not a callable.
+  """
+  if cluster is not None:
+    if isinstance(cluster, server_lib.ClusterSpec):
+      cluster_spec = cluster.as_dict()
+    else:
+      cluster_spec = server_lib.ClusterSpec(cluster).as_dict()
+    # Get ps_job_name from ps_device by stripping "/job:".
+    ps_job_name = pydev.DeviceSpec.from_string(ps_device).job
+    if ps_job_name not in cluster_spec or cluster_spec[ps_job_name] is None:
+      return None
+    ps_tasks = len(cluster_spec[ps_job_name])
+
+  if ps_tasks == 0:
+    return None
+
+  if ps_ops is None:
+    # TODO(sherrym): Variables in the LOCAL_VARIABLES collection should not be
+    # placed in the parameter server.
+    ps_ops = list(STANDARD_PS_OPS)
+
+  if not merge_devices:
+    logging.warning(
+        "DEPRECATION: It is recommended to set merge_devices=true in "
+        "replica_device_setter")
+  if ps_strategy is None:
+    ps_strategy = _RoundRobinStrategy(ps_tasks)
+  if not callable(ps_strategy):
+    raise TypeError("ps_strategy must be callable")
+  chooser = _ReplicaDeviceChooser(ps_tasks, ps_device, worker_device,
+                                  merge_devices, ps_ops, ps_strategy)
+  return chooser.device_function

tensorflow/python/training/device_setter_test.py

@@ -0,0 +1,187 @@
+# 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.
+# ==============================================================================
+"""Tests for device function for replicated training."""
+
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import device_setter
+from tensorflow.python.training import server_lib
+
+
+class DeviceSetterTest(test.TestCase):
+
+  _cluster_spec = server_lib.ClusterSpec({
+      "ps": ["ps0:2222", "ps1:2222"],
+      "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
+  })
+
+  @test_util.run_deprecated_v1
+  def testCPUOverride(self):
+    with ops.device(
+        device_setter.replica_device_setter(cluster=self._cluster_spec)):
+      with ops.device("/cpu:0"):
+        v = variables.Variable([1, 2])
+      w = variables.Variable([2, 1])
+      with ops.device("/cpu:0"):
+        a = v + w
+      self.assertDeviceEqual("/job:ps/task:0/cpu:0", v.device)
+      self.assertDeviceEqual("/job:ps/task:0/cpu:0", v.initializer.device)
+      self.assertDeviceEqual("/job:ps/task:1", w.device)
+      self.assertDeviceEqual("/job:ps/task:1", w.initializer.device)
+      self.assertDeviceEqual("/job:worker/cpu:0", a.device)
+
+  @test_util.run_deprecated_v1
+  def testResource(self):
+    with ops.device(
+        device_setter.replica_device_setter(cluster=self._cluster_spec)):
+      v = resource_variable_ops.ResourceVariable([1, 2])
+      self.assertDeviceEqual("/job:ps/task:0", v.device)
+
+  @test_util.run_deprecated_v1
+  def testPS2TasksWithClusterSpecClass(self):
+    with ops.device(
+        device_setter.replica_device_setter(cluster=self._cluster_spec)):
+      v = variables.Variable([1, 2])
+      w = variables.Variable([2, 1])
+      a = v + w
+      self.assertDeviceEqual("/job:ps/task:0", v.device)
+      self.assertDeviceEqual("/job:ps/task:0", v.initializer.device)
+      self.assertDeviceEqual("/job:ps/task:1", w.device)
+      self.assertDeviceEqual("/job:ps/task:1", w.initializer.device)
+      self.assertDeviceEqual("/job:worker", a.device)
+
+  @test_util.run_deprecated_v1
+  def testPS2TasksPinVariableToJob(self):
+    with ops.device(
+        device_setter.replica_device_setter(cluster=self._cluster_spec)):
+      v = variables.Variable([1, 2])
+      with ops.device("/job:moon"):
+        w = variables.Variable([2, 1])
+        with ops.device("/job:ps"):  # Explicit PS job will get task set.
+          x = variables.Variable([0, 1])
+      a = v + w + x
+      self.assertDeviceEqual("/job:ps/task:0", v.device)
+      self.assertDeviceEqual("/job:ps/task:0", v.initializer.device)
+      self.assertDeviceEqual("/job:moon", w.device)
+      self.assertDeviceEqual("/job:moon", w.initializer.device)
+      self.assertDeviceEqual("/job:ps/task:1", x.device)
+      self.assertDeviceEqual("/job:ps/task:1", x.initializer.device)
+      self.assertDeviceEqual("/job:worker", a.device)
+
+  @test_util.run_deprecated_v1
+  def testPS2TasksUseCpuForPS(self):
+    with ops.device(
+        device_setter.replica_device_setter(ps_tasks=1, ps_device="/cpu:0")):
+      v = variables.Variable([1, 2])
+      with ops.device("/job:moon"):
+        w = variables.Variable([2, 1])
+      a = v + w
+      self.assertDeviceEqual("/cpu:0", v.device)
+      self.assertDeviceEqual("/cpu:0", v.initializer.device)
+      self.assertDeviceEqual("/job:moon/cpu:0", w.device)
+      self.assertDeviceEqual("/job:moon/cpu:0", w.initializer.device)
+      self.assertDeviceEqual("/job:worker", a.device)
+
+  @test_util.run_deprecated_v1
+  def testPS2TasksNoMerging(self):
+    with ops.device(
+        device_setter.replica_device_setter(
+            cluster=self._cluster_spec, merge_devices=False)):
+      v = variables.Variable([1, 2])
+      with ops.device("/job:ps"):  # Won't assign task when merge_devices=False.
+        w = variables.Variable([2, 1])
+      a = v + w
+      self.assertDeviceEqual("/job:ps/task:0", v.device)
+      self.assertDeviceEqual("/job:ps/task:0", v.initializer.device)
+      self.assertDeviceEqual("/job:ps", w.device)
+      self.assertDeviceEqual("/job:ps", w.initializer.device)
+      self.assertDeviceEqual("/job:worker", a.device)
+
+  @test_util.run_deprecated_v1
+  def testPS2TasksWithClusterSpecDict(self):
+    with ops.device(
+        device_setter.replica_device_setter(cluster=self._cluster_spec.as_dict(
+        ))):
+      v = variables.Variable([1, 2])
+      w = variables.Variable([2, 1])
+      a = v + w
+      self.assertDeviceEqual("/job:ps/task:0", v.device)
+      self.assertDeviceEqual("/job:ps/task:0", v.initializer.device)
+      self.assertDeviceEqual("/job:ps/task:1", w.device)
+      self.assertDeviceEqual("/job:ps/task:1", w.initializer.device)
+      self.assertDeviceEqual("/job:worker", a.device)
+
+  @test_util.run_deprecated_v1
+  def testPS2TasksWithClusterDef(self):
+    with ops.device(
+        device_setter.replica_device_setter(
+            cluster=self._cluster_spec.as_cluster_def())):
+      v = variables.Variable([1, 2])
+      w = variables.Variable([2, 1])
+      a = v + w
+      self.assertDeviceEqual("/job:ps/task:0", v.device)
+      self.assertDeviceEqual("/job:ps/task:0", v.initializer.device)
+      self.assertDeviceEqual("/job:ps/task:1", w.device)
+      self.assertDeviceEqual("/job:ps/task:1", w.initializer.device)
+      self.assertDeviceEqual("/job:worker", a.device)
+
+  @test_util.run_deprecated_v1
+  def testPS2TasksWithDevice(self):
+    cluster_spec = server_lib.ClusterSpec({
+        "sun": ["sun0:2222", "sun1:2222", "sun2:2222"],
+        "moon": ["moon0:2222", "moon1:2222"]
+    })
+
+    with ops.device(
+        device_setter.replica_device_setter(
+            ps_device="/job:moon",
+            worker_device="/job:sun",
+            cluster=cluster_spec.as_cluster_def())):
+      v = variables.Variable([1, 2])
+      w = variables.Variable([2, 1])
+      a = v + w
+      self.assertDeviceEqual("/job:moon/task:0", v.device)
+      self.assertDeviceEqual("/job:moon/task:0", v.initializer.device)
+      self.assertDeviceEqual("/job:moon/task:1", w.device)
+      self.assertDeviceEqual("/job:moon/task:1", w.initializer.device)
+      self.assertDeviceEqual("/job:sun", a.device)
+
+  @test_util.run_deprecated_v1
+  def testPS2TasksWithCPUConstraint(self):
+    cluster_spec = server_lib.ClusterSpec({
+        "sun": ["sun0:2222", "sun1:2222", "sun2:2222"],
+        "moon": ["moon0:2222", "moon1:2222"]
+    })
+
+    with ops.device(
+        device_setter.replica_device_setter(
+            ps_device="/job:moon/cpu:0",
+            worker_device="/job:sun",
+            cluster=cluster_spec.as_cluster_def())):
+      v = variables.Variable([1, 2])
+      w = variables.Variable([2, 1])
+      a = v + w
+      self.assertDeviceEqual("/job:moon/task:0/cpu:0", v.device)
+      self.assertDeviceEqual("/job:moon/task:0/cpu:0", v.initializer.device)
+      self.assertDeviceEqual("/job:moon/task:1/cpu:0", w.device)
+      self.assertDeviceEqual("/job:moon/task:1/cpu:0", w.initializer.device)
+      self.assertDeviceEqual("/job:sun", a.device)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/evaluation.py

@@ -0,0 +1,273 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Contains functions for evaluation and summarization of metrics."""
+
+import math
+import time
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import state_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.training import basic_session_run_hooks
+from tensorflow.python.training import monitored_session
+from tensorflow.python.training import session_run_hook
+
+
+def _get_or_create_eval_step():
+  """Gets or creates the eval step `Tensor`.
+
+  Returns:
+    A `Tensor` representing a counter for the evaluation step.
+
+  Raises:
+    ValueError: If multiple `Tensors` have been added to the
+      `tf.GraphKeys.EVAL_STEP` collection.
+  """
+  graph = ops.get_default_graph()
+  eval_steps = graph.get_collection(ops.GraphKeys.EVAL_STEP)
+  if len(eval_steps) == 1:
+    return eval_steps[0]
+  elif len(eval_steps) > 1:
+    raise ValueError('Multiple tensors added to tf.GraphKeys.EVAL_STEP')
+  else:
+    counter = variable_scope.get_variable(
+        'eval_step',
+        shape=[],
+        dtype=dtypes.int64,
+        initializer=init_ops.zeros_initializer(),
+        trainable=False,
+        collections=[ops.GraphKeys.LOCAL_VARIABLES, ops.GraphKeys.EVAL_STEP])
+    return counter
+
+
+def _get_latest_eval_step_value(update_ops):
+  """Gets the eval step `Tensor` value after running `update_ops`.
+
+  Args:
+    update_ops: A list of `Tensors` or a dictionary of names to `Tensors`, which
+      are run before reading the eval step value.
+
+  Returns:
+    A `Tensor` representing the value for the evaluation step.
+  """
+  if isinstance(update_ops, dict):
+    update_ops = list(update_ops.values())
+
+  with ops.control_dependencies(update_ops):
+    return array_ops.identity(_get_or_create_eval_step().read_value())
+
+
+class _MultiStepStopAfterNEvalsHook(session_run_hook.SessionRunHook):
+  """Run hook used by the evaluation routines to run the `eval_ops` N times."""
+
+  def __init__(self, num_evals, steps_per_run=1):
+    """Constructs the run hook.
+
+    Args:
+      num_evals: The number of evaluations to run for. if set to None, will
+        iterate the dataset until all inputs are exhausted.
+      steps_per_run: Number of steps executed per run call.
+    """
+    self._num_evals = num_evals
+    self._evals_completed = None
+    self._steps_per_run_initial_value = steps_per_run
+
+  def _set_evals_completed_tensor(self, updated_eval_step):
+    self._evals_completed = updated_eval_step
+
+  def begin(self):
+    self._steps_per_run_variable = \
+        basic_session_run_hooks.get_or_create_steps_per_run_variable()
+
+  def after_create_session(self, session, coord):
+    # Update number of steps to run in the first run call
+    if self._num_evals is None:
+      steps = self._steps_per_run_initial_value
+    else:
+      steps = min(self._steps_per_run_initial_value, self._num_evals)
+    self._steps_per_run_variable.load(steps, session=session)
+
+  def before_run(self, run_context):
+    return session_run_hook.SessionRunArgs(
+        {'evals_completed': self._evals_completed})
+
+  def after_run(self, run_context, run_values):
+    evals_completed = run_values.results['evals_completed']
+    # Update number of steps to run in the next iteration
+    if self._num_evals is None:
+      steps = self._steps_per_run_initial_value
+    else:
+      steps = min(self._num_evals - evals_completed,
+                  self._steps_per_run_initial_value)
+    self._steps_per_run_variable.load(steps, session=run_context.session)
+
+    if self._num_evals is None:
+      logging.info('Evaluation [%d]', evals_completed)
+    else:
+      logging.info('Evaluation [%d/%d]', evals_completed, self._num_evals)
+    if self._num_evals is not None and evals_completed >= self._num_evals:
+      run_context.request_stop()
+
+
+class _StopAfterNEvalsHook(session_run_hook.SessionRunHook):
+  """Run hook used by the evaluation routines to run the `eval_ops` N times."""
+
+  def __init__(self, num_evals, log_progress=True):
+    """Constructs the run hook.
+
+    Args:
+      num_evals: The number of evaluations to run for. if set to None, will
+        iterate the dataset until all inputs are exhausted.
+      log_progress: Whether to log evaluation progress, defaults to True.
+    """
+    # The number of evals to run for.
+    self._num_evals = num_evals
+    self._evals_completed = None
+    self._log_progress = log_progress
+    # Reduce logging frequency if there are 20 or more evaluations.
+    self._log_frequency = (1 if (num_evals is None or num_evals < 20) else
+                           math.floor(num_evals / 10.))
+
+  def _set_evals_completed_tensor(self, updated_eval_step):
+    self._evals_completed = updated_eval_step
+
+  def before_run(self, run_context):
+    return session_run_hook.SessionRunArgs(
+        {'evals_completed': self._evals_completed})
+
+  def after_run(self, run_context, run_values):
+    evals_completed = run_values.results['evals_completed']
+    if self._log_progress:
+      if self._num_evals is None:
+        logging.info('Evaluation [%d]', evals_completed)
+      else:
+        if ((evals_completed % self._log_frequency) == 0 or
+            (self._num_evals == evals_completed)):
+          logging.info('Evaluation [%d/%d]', evals_completed, self._num_evals)
+    if self._num_evals is not None and evals_completed >= self._num_evals:
+      run_context.request_stop()
+
+
+def _evaluate_once(checkpoint_path,
+                   master='',
+                   scaffold=None,
+                   eval_ops=None,
+                   feed_dict=None,
+                   final_ops=None,
+                   final_ops_feed_dict=None,
+                   hooks=None,
+                   config=None):
+  """Evaluates the model at the given checkpoint path.
+
+  During a single evaluation, the `eval_ops` is run until the session is
+  interrupted or requested to finish. This is typically requested via a
+  `tf.contrib.training.StopAfterNEvalsHook` which results in `eval_ops` running
+  the requested number of times.
+
+  Optionally, a user can pass in `final_ops`, a single `Tensor`, a list of
+  `Tensors` or a dictionary from names to `Tensors`. The `final_ops` is
+  evaluated a single time after `eval_ops` has finished running and the fetched
+  values of `final_ops` are returned. If `final_ops` is left as `None`, then
+  `None` is returned.
+
+  One may also consider using a `tf.contrib.training.SummaryAtEndHook` to record
+  summaries after the `eval_ops` have run. If `eval_ops` is `None`, the
+  summaries run immediately after the model checkpoint has been restored.
+
+  Note that `evaluate_once` creates a local variable used to track the number of
+  evaluations run via `tf.contrib.training.get_or_create_eval_step`.
+  Consequently, if a custom local init op is provided via a `scaffold`, the
+  caller should ensure that the local init op also initializes the eval step.
+
+  Args:
+    checkpoint_path: The path to a checkpoint to use for evaluation.
+    master: The BNS address of the TensorFlow master.
+    scaffold: An tf.compat.v1.train.Scaffold instance for initializing variables
+      and restoring variables. Note that `scaffold.init_fn` is used by the
+      function to restore the checkpoint. If you supply a custom init_fn, then
+      it must also take care of restoring the model from its checkpoint.
+    eval_ops: A single `Tensor`, a list of `Tensors` or a dictionary of names to
+      `Tensors`, which is run until the session is requested to stop, commonly
+      done by a `tf.contrib.training.StopAfterNEvalsHook`.
+    feed_dict: The feed dictionary to use when executing the `eval_ops`.
+    final_ops: A single `Tensor`, a list of `Tensors` or a dictionary of names
+      to `Tensors`.
+    final_ops_feed_dict: A feed dictionary to use when evaluating `final_ops`.
+    hooks: List of `tf.estimator.SessionRunHook` callbacks which are run inside
+      the evaluation loop.
+    config: An instance of `tf.compat.v1.ConfigProto` that will be used to
+      configure the `Session`. If left as `None`, the default will be used.
+
+  Returns:
+    The fetched values of `final_ops` or `None` if `final_ops` is `None`.
+  """
+  eval_step = _get_or_create_eval_step()
+
+  # Prepare the run hooks.
+  hooks = list(hooks or [])
+
+  if eval_ops is not None:
+    if any(isinstance(h, _MultiStepStopAfterNEvalsHook) for h in hooks):
+      steps_per_run_variable = \
+          basic_session_run_hooks.get_or_create_steps_per_run_variable()
+      update_eval_step = state_ops.assign_add(
+          eval_step,
+          math_ops.cast(steps_per_run_variable, dtype=eval_step.dtype),
+          use_locking=True)
+    else:
+      update_eval_step = state_ops.assign_add(eval_step, 1, use_locking=True)
+
+    if isinstance(eval_ops, dict):
+      eval_ops['update_eval_step'] = update_eval_step
+    elif isinstance(eval_ops, (tuple, list)):
+      eval_ops = list(eval_ops) + [update_eval_step]
+    else:
+      eval_ops = [eval_ops, update_eval_step]
+
+    eval_step_value = _get_latest_eval_step_value(eval_ops)
+
+    for h in hooks:
+      if isinstance(h, (_StopAfterNEvalsHook, _MultiStepStopAfterNEvalsHook)):
+        h._set_evals_completed_tensor(eval_step_value)  # pylint: disable=protected-access
+
+  logging.info('Starting evaluation at ' +
+               time.strftime('%Y-%m-%dT%H:%M:%S', time.localtime()))
+  start = time.time()
+  # Prepare the session creator.
+  session_creator = monitored_session.ChiefSessionCreator(
+      scaffold=scaffold,
+      checkpoint_filename_with_path=checkpoint_path,
+      master=master,
+      config=config)
+
+  final_ops_hook = basic_session_run_hooks.FinalOpsHook(final_ops,
+                                                        final_ops_feed_dict)
+  hooks.append(final_ops_hook)
+
+  with monitored_session.MonitoredSession(
+      session_creator=session_creator, hooks=hooks) as session:
+    if eval_ops is not None:
+      while not session.should_stop():
+        session.run(eval_ops, feed_dict)
+  logging.info('Inference Time : {:0.5f}s'.format(time.time() - start))
+
+  logging.info('Finished evaluation at ' +
+               time.strftime('%Y-%m-%d-%H:%M:%S', time.localtime()))
+  return final_ops_hook.final_ops_values

tensorflow/python/training/evaluation_test.py

@@ -0,0 +1,247 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Tests for tf.training.evaluation."""
+
+import os
+
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import random_seed
+from tensorflow.python.layers import layers
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import metrics as metrics_module
+from tensorflow.python.ops import state_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops.losses import losses
+from tensorflow.python.platform import test
+from tensorflow.python.training import basic_session_run_hooks
+from tensorflow.python.training import evaluation
+from tensorflow.python.training import gradient_descent
+from tensorflow.python.training import monitored_session
+from tensorflow.python.training import saver
+from tensorflow.python.training import training
+
+_USE_GLOBAL_STEP = 0
+
+
+def logistic_classifier(inputs):
+  return layers.dense(inputs, 1, activation=math_ops.sigmoid)
+
+
+def local_variable(init_value, name):
+  return variable_scope.get_variable(
+      name,
+      dtype=dtypes.float32,
+      initializer=init_value,
+      trainable=False,
+      collections=[ops.GraphKeys.LOCAL_VARIABLES])
+
+
+class EvaluateOnceTest(test.TestCase):
+
+  def setUp(self):
+    super(EvaluateOnceTest, self).setUp()
+
+    # Create an easy training set:
+    np.random.seed(0)
+
+    self._inputs = np.zeros((16, 4))
+    self._labels = np.random.randint(0, 2, size=(16, 1)).astype(np.float32)
+
+    for i in range(16):
+      j = int(2 * self._labels[i] + np.random.randint(0, 2))
+      self._inputs[i, j] = 1
+
+  def _train_model(self, checkpoint_dir, num_steps):
+    """Trains a simple classification model.
+
+    Note that the data has been configured such that after around 300 steps,
+    the model has memorized the dataset (e.g. we can expect %100 accuracy).
+
+    Args:
+      checkpoint_dir: The directory where the checkpoint is written to.
+      num_steps: The number of steps to train for.
+    """
+    with ops.Graph().as_default():
+      random_seed.set_random_seed(0)
+      tf_inputs = constant_op.constant(self._inputs, dtype=dtypes.float32)
+      tf_labels = constant_op.constant(self._labels, dtype=dtypes.float32)
+
+      tf_predictions = logistic_classifier(tf_inputs)
+      loss_op = losses.log_loss(labels=tf_labels, predictions=tf_predictions)
+
+      optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=1.0)
+      train_op = optimizer.minimize(loss_op,
+                                    training.get_or_create_global_step())
+
+      with monitored_session.MonitoredTrainingSession(
+          checkpoint_dir=checkpoint_dir,
+          hooks=[basic_session_run_hooks.StopAtStepHook(num_steps)]) as session:
+        loss = None
+        while not session.should_stop():
+          _, loss = session.run([train_op, loss_op])
+
+        if num_steps >= 300:
+          assert loss < .015
+
+  def testEvaluatePerfectModel(self):
+    checkpoint_dir = os.path.join(self.get_temp_dir(),
+                                  'evaluate_perfect_model_once')
+
+    # Train a Model to completion:
+    self._train_model(checkpoint_dir, num_steps=300)
+
+    # Run
+    inputs = constant_op.constant(self._inputs, dtype=dtypes.float32)
+    labels = constant_op.constant(self._labels, dtype=dtypes.float32)
+    logits = logistic_classifier(inputs)
+    predictions = math_ops.round(logits)
+
+    accuracy, update_op = metrics_module.accuracy(labels, predictions)
+
+    checkpoint_path = saver.latest_checkpoint(checkpoint_dir)
+
+    final_ops_values = evaluation._evaluate_once(
+        checkpoint_path=checkpoint_path,
+        eval_ops=update_op,
+        final_ops={'accuracy': (accuracy, update_op)},
+        hooks=[
+            evaluation._StopAfterNEvalsHook(1),
+        ])
+    self.assertGreater(final_ops_values['accuracy'], .99)
+
+  def testEvaluateWithFiniteInputs(self):
+    checkpoint_dir = os.path.join(self.get_temp_dir(),
+                                  'evaluate_with_finite_inputs')
+
+    # Train a Model to completion:
+    self._train_model(checkpoint_dir, num_steps=300)
+
+    # Run evaluation. Inputs are fed through input producer for one epoch.
+    all_inputs = constant_op.constant(self._inputs, dtype=dtypes.float32)
+    all_labels = constant_op.constant(self._labels, dtype=dtypes.float32)
+
+    single_input, single_label = training.slice_input_producer(
+        [all_inputs, all_labels], num_epochs=1)
+    inputs, labels = training.batch([single_input, single_label], batch_size=6,
+                                    allow_smaller_final_batch=True)
+
+    logits = logistic_classifier(inputs)
+    predictions = math_ops.round(logits)
+
+    accuracy, update_op = metrics_module.accuracy(labels, predictions)
+
+    checkpoint_path = saver.latest_checkpoint(checkpoint_dir)
+
+    final_ops_values = evaluation._evaluate_once(
+        checkpoint_path=checkpoint_path,
+        eval_ops=update_op,
+        final_ops={
+            'accuracy': (accuracy, update_op),
+            'eval_steps': evaluation._get_or_create_eval_step()
+        },
+        hooks=[
+            evaluation._StopAfterNEvalsHook(None),
+        ])
+    self.assertTrue(final_ops_values['accuracy'] > .99)
+    # Runs evaluation for 4 iterations. First 2 evaluate full batch of 6 inputs
+    # each; the 3rd iter evaluates the remaining 4 inputs, and the last one
+    # triggers an error which stops evaluation.
+    self.assertEqual(final_ops_values['eval_steps'], 4)
+
+  def testEvalOpAndFinalOp(self):
+    checkpoint_dir = os.path.join(self.get_temp_dir(), 'eval_ops_and_final_ops')
+
+    # Train a model for a single step to get a checkpoint.
+    self._train_model(checkpoint_dir, num_steps=1)
+    checkpoint_path = saver.latest_checkpoint(checkpoint_dir)
+
+    # Create the model so we have something to restore.
+    inputs = constant_op.constant(self._inputs, dtype=dtypes.float32)
+    logistic_classifier(inputs)
+
+    num_evals = 5
+    final_increment = 9.0
+
+    my_var = local_variable(0.0, name='MyVar')
+    eval_ops = state_ops.assign_add(my_var, 1.0)
+    final_ops = array_ops.identity(my_var) + final_increment
+
+    final_hooks = [evaluation._StopAfterNEvalsHook(num_evals),]
+    initial_hooks = list(final_hooks)
+    final_ops_values = evaluation._evaluate_once(
+        checkpoint_path=checkpoint_path,
+        eval_ops=eval_ops,
+        final_ops={'value': final_ops},
+        hooks=final_hooks)
+    self.assertEqual(final_ops_values['value'], num_evals + final_increment)
+    self.assertEqual(initial_hooks, final_hooks)
+
+  def testMultiEvalStepIncrements(self):
+    checkpoint_dir = os.path.join(self.get_temp_dir(), 'eval_ops_and_final_ops')
+
+    # Train a model for a single step to get a checkpoint.
+    self._train_model(checkpoint_dir, num_steps=1)
+    checkpoint_path = saver.latest_checkpoint(checkpoint_dir)
+
+    # Create the model so we have something to restore.
+    inputs = constant_op.constant(self._inputs, dtype=dtypes.float32)
+    logistic_classifier(inputs)
+
+    num_evals = 6
+
+    my_var = local_variable(0.0, name='MyVar')
+    # In eval ops, we also increase the eval step one more time.
+    eval_ops = [state_ops.assign_add(my_var, 1.0),
+                state_ops.assign_add(
+                    evaluation._get_or_create_eval_step(), 1, use_locking=True)]
+    expect_eval_update_counts = num_evals // 2
+
+    final_ops = array_ops.identity(my_var)
+
+    final_ops_values = evaluation._evaluate_once(
+        checkpoint_path=checkpoint_path,
+        eval_ops=eval_ops,
+        final_ops={'value': final_ops},
+        hooks=[evaluation._StopAfterNEvalsHook(num_evals),])
+    self.assertEqual(final_ops_values['value'], expect_eval_update_counts)
+
+  def testOnlyFinalOp(self):
+    checkpoint_dir = os.path.join(self.get_temp_dir(), 'only_final_ops')
+
+    # Train a model for a single step to get a checkpoint.
+    self._train_model(checkpoint_dir, num_steps=1)
+    checkpoint_path = saver.latest_checkpoint(checkpoint_dir)
+
+    # Create the model so we have something to restore.
+    inputs = constant_op.constant(self._inputs, dtype=dtypes.float32)
+    logistic_classifier(inputs)
+
+    final_increment = 9.0
+
+    my_var = local_variable(0.0, name='MyVar')
+    final_ops = array_ops.identity(my_var) + final_increment
+
+    final_ops_values = evaluation._evaluate_once(
+        checkpoint_path=checkpoint_path, final_ops={'value': final_ops})
+    self.assertEqual(final_ops_values['value'], final_increment)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/training/experimental/BUILD

@@ -0,0 +1,148 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow:internal"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "loss_scale",
+    srcs = ["loss_scale.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/distribute:reduce_util",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:indexed_slices",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:control_flow_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:variable_v1",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/trackable:base",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "loss_scale_optimizer",
+    srcs = ["loss_scale_optimizer.py"],
+    strict_deps = True,
+    deps = [
+        ":loss_scale",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/framework:indexed_slices",
+        "//tensorflow/python/framework:smart_cond",
+        "//tensorflow/python/ops:control_flow_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/training:optimizer",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_test(
+    name = "loss_scale_optimizer_test",
+    size = "small",
+    srcs = ["loss_scale_optimizer_test.py"],
+    strict_deps = True,
+    deps = [
+        ":loss_scale",
+        ":loss_scale_optimizer",
+        "//tensorflow/python/checkpoint",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/distribute:mirrored_strategy",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_conversion",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:check_ops",
+        "//tensorflow/python/ops:custom_gradient",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/training:gradient_descent",
+        "//tensorflow/python/training:momentum",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_test(
+    name = "loss_scale_test",
+    size = "medium",
+    srcs = ["loss_scale_test.py"],
+    strict_deps = True,
+    deps = [
+        ":loss_scale",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/distribute:mirrored_strategy",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:check_ops",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_library(
+    name = "mixed_precision_global_state",
+    srcs = ["mixed_precision_global_state.py"],
+    strict_deps = True,
+    deps = ["//tensorflow/python/util:tf_export"],
+)
+
+py_library(
+    name = "mixed_precision",
+    srcs = ["mixed_precision.py"],
+    strict_deps = True,
+    deps = [
+        ":loss_scale_optimizer",
+        ":mixed_precision_global_state",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/training:optimizer",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "mixed_precision_test",
+    size = "small",
+    srcs = ["mixed_precision_test.py"],
+    deps = [
+        ":loss_scale_optimizer",
+        ":mixed_precision",
+        ":mixed_precision_global_state",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/client:session",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/training:gradient_descent",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)

tensorflow/python/training/experimental/loss_scale.py

@@ -0,0 +1,453 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Contains LossScale classes."""
+import abc
+
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import reduce_util
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.ops import variables
+from tensorflow.python.trackable import base as trackable
+from tensorflow.python.util import deprecation
+from tensorflow.python.util import nest
+from tensorflow.python.util.tf_export import tf_export
+
+
+@deprecation.deprecated_endpoints('mixed_precision.experimental.LossScale',
+                                  'train.experimental.LossScale')
+@tf_export(
+    v1=[
+        'mixed_precision.LossScale',
+        'mixed_precision.experimental.LossScale',
+        'train.experimental.LossScale'
+    ])
+class LossScale(trackable.Trackable, metaclass=abc.ABCMeta):
+  """Base class for all TF1 loss scales.
+
+  This is an abstract base class, so you cannot instantiate it directly.
+  Instead, use one of its concrete subclasses:
+    * `tf.compat.v1.mixed_precision.DynamicLossScale`
+    * `tf.compat.v1.mixed_precision.FixedLossScale`
+
+  Loss scaling is a process that multiplies the loss by a multiplier called the
+  loss scale, and divides each gradient by the same multiplier. The pseudocode
+  for this process is:
+
+  ```
+  loss = ...
+  loss *= loss_scale
+  grads = gradients(loss, vars)
+  grads /= loss_scale
+  ```
+
+  Mathematically, loss scaling has no effect, but can help avoid numerical
+  underflow in intermediate gradients when float16 tensors are used for mixed
+  precision training. By multiplying the loss, each intermediate gradient will
+  have the same multiplier applied.
+
+  Instances of this class represent a loss scale. Calling instances of this
+  class returns the loss scale as a scalar float32 tensor, while method
+  `update()` updates the loss scale depending on the values of the gradients.
+  Optimizers use instances of this class to scale loss and gradients.
+
+  In most functions that accept a LossScale, you can also pass an int (such as
+  8) to create a `FixedLossScale` or the string `"dynamic"` to create a dynamic
+  loss scale.
+  """
+
+  def __init__(self):
+    """Initializes the loss scale class."""
+    self._weights = {}
+
+  @abc.abstractmethod
+  def __call__(self):
+    """Returns the current loss scale as a scalar `float32` tensor."""
+    pass
+
+  @abc.abstractmethod
+  def update(self, grads):
+    """Updates the value of the loss scale.
+
+    The loss scale will be potentially updated, based on the value of `grads`.
+    The tensor returned by calling this class is only updated when this function
+    is evaluated.
+
+    In eager mode, this directly updates the loss scale, so that calling
+    `__call__` will return the newly updated loss scale. In graph mode,
+    this returns an op that, when evaluated, updates the loss scale.
+
+    This function also returns a `should_apply_gradients` bool. If False,
+    gradients should not be applied to the variables that step, as nonfinite
+    gradients were found, and the loss scale has been be updated to reduce the
+    chance of finding nonfinite gradients in the next step. Some loss scale
+    classes will always return True, as they cannot adjust themselves in
+    response to nonfinite gradients.
+
+    When a DistributionStrategy is used, this function may only be called in a
+    cross-replica context.
+
+    Args:
+      grads: A nested structure of unscaled gradients, each which is the
+        gradient of the loss with respect to a weight. The gradients should have
+        already been divided by the loss scale being before passed to this
+        function. 'None' gradients are accepted, and are ignored.
+
+    Returns:
+      update_op: In eager mode, None. In graph mode, an op to update the loss
+        scale.
+      should_apply_gradients: Either a bool or a scalar boolean tensor. If
+        False, the caller should skip applying `grads` to the variables this
+        step.
+    """
+    pass
+
+  def _add_weight(self, name, initial_value, dtype=None):
+    """Adds a weight to this loss scale.
+
+    Args:
+      name: Variable name.
+      initial_value: The variable's initial value.
+      dtype: The type of the variable.
+
+    Returns:
+      A variable.
+
+    Raises:
+      RuntimeError: If a weight with `name` has already been added.
+    """
+    variable = variable_v1.VariableV1(
+        initial_value=initial_value,
+        name=name,
+        dtype=dtype,
+        trainable=False,
+        use_resource=True,
+        synchronization=variables.VariableSynchronization.AUTO,
+        # Set aggregation to NONE, as loss scaling variables should never be
+        # aggregated.
+        aggregation=variables.VariableAggregation.NONE)
+    if context.executing_eagerly():
+      graph_key = None
+    else:
+      graph = ops.get_default_graph()
+      graph_key = graph._graph_key  # pylint: disable=protected-access
+
+    key = (name, graph_key)
+    if self._weights.get(key, None) is not None:
+      raise RuntimeError('Duplicate variables detected. {}'.format(key))
+    self._weights[key] = variable
+    self._handle_deferred_dependencies(name=name, trackable=variable)
+    return variable
+
+  def _trackable_children(self,
+                          save_type=trackable.SaveType.CHECKPOINT,
+                          **kwargs):
+    """From Trackable. Gather graph-specific weights to save."""
+    if context.executing_eagerly():
+      graph_key = None
+    else:
+      graph = ops.get_default_graph()
+      graph_key = graph._graph_key  # pylint: disable=protected-access
+    weights = {}
+    for (name, g), v in sorted(self._weights.items(), key=lambda i: i[0][0]):
+      if g == graph_key:
+        weights[name] = v
+    weights.update(
+        super(LossScale, self)._trackable_children(save_type, **kwargs))
+    return weights
+
+  def _lookup_dependency(self, name, cached_dependencies=None):
+    """From Trackable. Find a weight in the current graph."""
+    unconditional = super(LossScale, self)._lookup_dependency(
+        name, cached_dependencies)
+    if unconditional is not None:
+      return unconditional
+    if context.executing_eagerly():
+      graph_key = None
+    else:
+      graph = ops.get_default_graph()
+      graph_key = graph._graph_key  # pylint: disable=protected-access
+    return self._weights.get((name, graph_key), None)
+
+  @abc.abstractmethod
+  def get_config(self):
+    """Returns the config of this loss scale."""
+    pass
+
+  @classmethod
+  def from_config(cls, config):
+    """Creates the LossScale from its config."""
+    return cls(**config)
+
+
+@deprecation.deprecated_endpoints('mixed_precision.experimental.FixedLossScale',
+                                  'train.experimental.FixedLossScale')
+@tf_export(
+    v1=[
+        'mixed_precision.FixedLossScale',
+        'mixed_precision.experimental.FixedLossScale',
+        'train.experimental.FixedLossScale'
+    ])
+class FixedLossScale(LossScale):
+  """Loss scale with a fixed value.
+
+  The loss scale is not updated for the lifetime of instances of this class.
+  A given instance of this class always returns the same number when called.
+  """
+
+  @deprecation.deprecated(
+      None, 'Use tf.keras.mixed_precision.LossScaleOptimizer instead. '
+            'LossScaleOptimizer now has all the functionality of '
+            'FixedLossScale')
+  def __init__(self, loss_scale_value):
+    """Creates the fixed loss scale.
+
+    Args:
+      loss_scale_value: A Python float. Its ideal value varies depending on
+        models to run. Choosing a too small loss_scale might affect model
+        quality; a too big loss_scale might cause inf or nan. There is no single
+        right loss_scale to apply. There is no harm choosing a relatively big
+        number as long as no nan or inf is encountered in training.
+
+    Raises:
+      ValueError: If loss_scale_value is less than 1.
+    """
+    super(FixedLossScale, self).__init__()
+    if not isinstance(loss_scale_value, (int, float)):
+      raise ValueError('loss_scale_value must be a Python int or float.')
+    if loss_scale_value < 1:
+      raise ValueError('loss_scale_value must be at least 1.')
+    # It's important we do not create tensors in the constructor, as such
+    # tensors might be on a different device or tf.function vs when the tensor
+    # is used. This would hurt performance. Therefore, we do not create a tensor
+    # from loss_scale_value, but instead leave it as a Python float.
+    # TODO(reedwm): Also do not create tensors in the DynamicLossScale
+    # constructor.
+    self._loss_scale_value = float(loss_scale_value)
+
+  def __call__(self):
+    return ops.convert_to_tensor(self._loss_scale_value)
+
+  def update(self, grads):
+    del grads
+    return control_flow_ops.no_op(), True
+
+  def __repr__(self):
+    return 'FixedLossScale(%s)' % self._loss_scale_value
+
+  def get_config(self):
+    return {'loss_scale_value': self._loss_scale_value}
+
+
+def _is_all_finite(grads):
+  """Returns a scalar boolean tensor indicating if all gradients are finite."""
+  def raw_values(g):
+    return g.values if isinstance(g, indexed_slices.IndexedSlices) else g
+
+  is_finite_per_grad = [
+      math_ops.reduce_all(math_ops.is_finite(raw_values(g)))
+      for g in grads
+      if g is not None
+  ]
+  return math_ops.reduce_all(is_finite_per_grad)
+
+
+def _op_in_graph_mode(tensor):
+  """Returns the tensor's op in graph mode, or the tensor in eager mode.
+
+  This is useful because sometimes an op is needed in graph mode instead of a
+  tensor. In eager mode, there are no ops.
+
+  Args:
+    tensor: A tensor.
+
+  Returns:
+    The tensor's op in graph mode. The tensor in eager mode.
+  """
+  if context.executing_eagerly():
+    return tensor
+  return tensor.op
+
+
+def _assign_if_finite(var, value):
+  """Assigns a value to a variable if the value is finite."""
+  return cond.cond(
+      math_ops.is_finite(value), lambda: _op_in_graph_mode(var.assign(value)),
+      control_flow_ops.no_op)
+
+
+@deprecation.deprecated_endpoints(
+    'mixed_precision.experimental.DynamicLossScale',
+    'train.experimental.DynamicLossScale')
+@tf_export(
+    v1=[
+        'mixed_precision.DynamicLossScale',
+        'mixed_precision.experimental.DynamicLossScale',
+        'train.experimental.DynamicLossScale'
+    ])
+class DynamicLossScale(LossScale):
+  """Loss scale that dynamically adjusts itself.
+
+  Dynamic loss scaling works by adjusting the loss scale as training progresses.
+  The goal is to keep the loss scale as high as possible without overflowing the
+  gradients. As long as the gradients do not overflow, raising the loss scale
+  never hurts.
+
+  The algorithm starts by setting the loss scale to an initial value. Every N
+  steps that the gradients are finite, the loss scale is increased by some
+  factor. However, if a NaN or Inf gradient is found, the gradients for that
+  step are not applied, and the loss scale is decreased by the factor. This
+  process tends to keep the loss scale as high as possible without gradients
+  overflowing.
+  """
+
+  @deprecation.deprecated(
+      None, 'Use tf.keras.mixed_precision.LossScaleOptimizer instead. '
+            'LossScaleOptimizer now has all the functionality of '
+            'DynamicLossScale')
+  def __init__(self,
+               initial_loss_scale=2 ** 15,  # See docstring for why this is big.
+               increment_period=2000,
+               multiplier=2.):
+    """Creates the dynamic loss scale.
+
+    Args:
+      initial_loss_scale: A Python float.  The loss scale to use at the
+        beginning. It's better to start this at a very high number, because a
+        loss scale that is too high gets lowered far more quickly than a loss
+        scale that is too low gets raised. The default is 2 ** 15, which is
+        approximately half the maximum float16 value.
+      increment_period: Increases loss scale every `increment_period`
+        consecutive steps that finite gradients are encountered. If a nonfinite
+        gradient is encountered, the count is reset back to zero.
+      multiplier: The multiplier to use when increasing or decreasing the loss
+        scale.
+    """
+    super(DynamicLossScale, self).__init__()
+    self._initial_loss_scale = float(initial_loss_scale)
+    self._increment_period = int(increment_period)
+    self._multiplier = float(multiplier)
+
+    self._current_loss_scale = self._add_weight(
+        name='current_loss_scale',
+        dtype=dtypes.float32,
+        initial_value=self._initial_loss_scale)
+    # The number of consecutive steps with finite gradients since the last
+    # nonfinite gradient or change in loss scale.
+    self._num_good_steps = self._add_weight(
+        name='good_steps', dtype=dtypes.int64, initial_value=0)
+
+  @property
+  def initial_loss_scale(self):
+    return self._initial_loss_scale
+
+  @property
+  def increment_period(self):
+    return self._increment_period
+
+  @property
+  def multiplier(self):
+    return self._multiplier
+
+  def __call__(self):
+    return ops.convert_to_tensor(self._current_loss_scale)
+
+  def update(self, grads):
+    """Updates loss scale based on if gradients are finite in current step."""
+    grads = nest.flatten(grads)
+    if distribute_lib.has_strategy():
+      distribution = distribute_lib.get_cross_replica_context()
+
+      def get_is_finite(grads):
+        is_finite = _is_all_finite(grads)
+        # We cast to float, because we cannot reduce booleans with
+        # DistributionStrategy.
+        return math_ops.cast(is_finite, dtypes.float32)
+
+      is_finite_float = distribution.extended.call_for_each_replica(
+          get_is_finite, args=(grads,))
+      reduced_is_finite_float = distribution.reduce(reduce_util.ReduceOp.SUM,
+                                                    is_finite_float, axis=None)
+      is_finite = math_ops.equal(reduced_is_finite_float,
+                                 distribution.num_replicas_in_sync)
+    else:
+      is_finite = _is_all_finite(grads)
+
+    def update_if_finite_grads():
+      """Update assuming the gradients are finite."""
+
+      def incr_loss_scale():
+        new_loss_scale = self._current_loss_scale * self._multiplier
+        return control_flow_ops.group(
+            _assign_if_finite(self._current_loss_scale, new_loss_scale),
+            self._num_good_steps.assign(0))
+
+      return cond.cond(
+          self._num_good_steps + 1 >= self._increment_period,
+          incr_loss_scale, lambda: _op_in_graph_mode(
+              self._num_good_steps.assign_add(1)))
+
+    def update_if_not_finite_grads():
+      """Update assuming the gradients are nonfinite."""
+
+      new_loss_scale = math_ops.maximum(
+          self._current_loss_scale / self._multiplier, 1)
+      return control_flow_ops.group(
+          self._num_good_steps.assign(0),
+          self._current_loss_scale.assign(new_loss_scale))
+
+    update_op = cond.cond(is_finite, update_if_finite_grads,
+                          update_if_not_finite_grads)
+    should_apply_gradients = is_finite
+    return update_op, should_apply_gradients
+
+  def __repr__(self):
+    if context.executing_eagerly():
+      return ('DynamicLossScale(current_loss_scale=%s, num_good_steps=%s, '
+              'initial_loss_scale=%s, increment_period=%s, multiplier=%s)' %
+              (self._current_loss_scale.numpy(), self._num_good_steps.numpy(),
+               self.initial_loss_scale, self.increment_period, self.multiplier))
+    else:
+      return ('DynamicLossScale(initial_loss_scale=%s, increment_period=%s, '
+              'multiplier=%s)' %
+              (self.initial_loss_scale, self.increment_period, self.multiplier))
+
+  def get_config(self):
+    return {
+        'initial_loss_scale': self.initial_loss_scale,
+        'increment_period': self.increment_period,
+        'multiplier': self.multiplier,
+    }
+
+
+def get(identifier):
+  """Get a loss scale object."""
+  if isinstance(identifier, (int, float)):
+    return FixedLossScale(identifier)
+  if identifier == 'dynamic':
+    return DynamicLossScale()
+  if isinstance(identifier, LossScale):
+    return identifier
+  elif identifier is None:
+    return None
+  else:
+    raise ValueError('Could not interpret loss scale identifier: %s' %
+                     identifier)

tensorflow/python/training/experimental/loss_scale_optimizer.py

@@ -0,0 +1,251 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Contains LossScale classes."""
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import smart_cond
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.training import optimizer
+from tensorflow.python.training.experimental import loss_scale as loss_scale_module
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+@deprecation.deprecated_endpoints(
+    'train.experimental.MixedPrecisionLossScaleOptimizer')
+@tf_export(v1=['mixed_precision.MixedPrecisionLossScaleOptimizer',
+               'train.experimental.MixedPrecisionLossScaleOptimizer'])
+class MixedPrecisionLossScaleOptimizer(optimizer.Optimizer):
+  """An optimizer that applies loss scaling.
+
+  Loss scaling is a process that multiplies the loss by a multiplier called the
+  loss scale, and divides each gradient by the same multiplier. The pseudocode
+  for this process is:
+
+  ```
+  loss = ...
+  loss *= loss_scale
+  grads = gradients(loss, vars)
+  grads /= loss_scale
+  ```
+
+  Mathematically, loss scaling has no effect, but can help avoid numerical
+  underflow in intermediate gradients when float16 tensors are used for mixed
+  precision training. By multiplying the loss, each intermediate gradient will
+  have the same multiplier applied.
+
+  The loss scale can either be a fixed constant, chosen by the user, or be
+  dynamically determined. Dynamically determining the loss scale is convenient
+  as a loss scale does not have to be explicitly chosen. However it reduces
+  performance.
+
+  This optimizer wraps another optimizer and applies loss scaling to it via a
+  `LossScale`. Loss scaling is applied whenever gradients are
+  computed, such as through `minimize()`.
+  """
+
+  def __init__(self, opt, loss_scale):
+    if not isinstance(opt, optimizer.Optimizer):
+      raise ValueError('"opt" must be an instance of Optimizer, but got: %s' %
+                       type(opt))
+    self._optimizer = opt
+
+    use_locking = opt._use_locking  # pylint: disable=protected-access
+    name = opt.get_name()
+    super(MixedPrecisionLossScaleOptimizer, self).__init__(use_locking, name)
+
+    self._loss_scale = loss_scale_module.get(loss_scale)
+    if self._loss_scale is None:
+      raise ValueError('loss_scale cannot be None')
+    self._track_trackable(self._optimizer, 'base_optimizer')
+    self._track_trackable(self._loss_scale, 'loss_scale')
+
+  def _doing_dynamic_loss_scaling(self):
+    """Check if `_loss_scale` dynamically manages the loss scale."""
+    return isinstance(self._loss_scale, loss_scale_module.DynamicLossScale)
+
+  def compute_gradients(self,
+                        loss,
+                        var_list=None,
+                        gate_gradients=optimizer.Optimizer.GATE_OP,
+                        aggregation_method=None,
+                        colocate_gradients_with_ops=False,
+                        grad_loss=None):
+    """Compute gradients of `loss` for the variables in `var_list`.
+
+    This adjusts the dynamic range of the gradient evaluation by scaling up
+    the `loss` value. The gradient values are then scaled back down by the
+    reciprocal of the loss scale. This is useful in reduced precision training
+    where small gradient values would otherwise underflow the representable
+    range.
+
+    Args:
+      loss: A Tensor containing the value to minimize or a callable taking no
+        arguments which returns the value to minimize. When eager execution is
+        enabled it must be a callable.
+      var_list: Optional list or tuple of `tf.Variable` to update to minimize
+        `loss`.  Defaults to the list of variables collected in the graph under
+        the key `GraphKeys.TRAINABLE_VARIABLES`.
+      gate_gradients: How to gate the computation of gradients.  Can be
+        `GATE_NONE`, `GATE_OP`, or `GATE_GRAPH`.
+      aggregation_method: Specifies the method used to combine gradient terms.
+        Valid values are defined in the class `AggregationMethod`.
+      colocate_gradients_with_ops: If True, try colocating gradients with the
+        corresponding op.
+      grad_loss: Optional. A `Tensor` holding the gradient computed for `loss`.
+
+    Returns:
+      A list of (gradient, variable) pairs. Variable is always present, but
+      gradient can be `None`.
+    """
+    loss = self._scale_loss(loss)
+    grads_and_vars = self._optimizer.compute_gradients(
+        loss=loss,
+        var_list=var_list,
+        gate_gradients=gate_gradients,
+        aggregation_method=aggregation_method,
+        colocate_gradients_with_ops=colocate_gradients_with_ops,
+        grad_loss=grad_loss)
+
+    grads = [g for g, _ in grads_and_vars]
+    variables = [v for _, v in grads_and_vars]
+    unscaled_grads = self._unscale_grads(grads)
+    return list(zip(unscaled_grads, variables))
+
+  def _scale_loss(self, loss):
+    loss_scale = self._loss_scale()
+    if callable(loss):
+      def new_loss():
+        loss_val = loss()
+        return loss_val * math_ops.cast(loss_scale, loss_val.dtype)
+      return new_loss
+    else:
+      return loss * math_ops.cast(loss_scale, loss.dtype)
+
+  def _unscale_grads(self, grads):
+    loss_scale = self._loss_scale()
+    loss_scale_reciprocal = 1 / loss_scale
+    return [
+        None if g is None else self._scale_grad(g, loss_scale_reciprocal)
+        for g in grads
+    ]
+
+  def _scale_grad(self, grad, loss_scale_reciprocal):
+    if isinstance(grad, indexed_slices.IndexedSlices):
+      grad_vals = grad.values * loss_scale_reciprocal
+      return indexed_slices.IndexedSlices(grad_vals, grad.indices,
+                                          grad.dense_shape)
+    return grad * loss_scale_reciprocal
+
+  def apply_gradients(self, grads_and_vars, global_step=None, name=None):
+    """Apply gradients to variables.
+
+    This is the second part of `minimize()`. It returns an `Operation` that
+    conditionally applies gradients if all gradient values are finite.
+    Otherwise no update is performed (nor is `global_step` incremented).
+
+    Args:
+      grads_and_vars: List of (gradient, variable) pairs as returned by
+        `compute_gradients()`.
+      global_step: Optional `Variable` to increment by one after the variables
+        have been updated.
+      name: Optional name for the returned operation.  Default to the name
+        passed to the `Optimizer` constructor.
+
+    Returns:
+      An `Operation` that conditionally applies the specified gradients. If
+      `global_step` was not None, that operation also increments `global_step`.
+
+    Raises:
+      RuntimeError: If you should use `_distributed_apply()` instead.
+    """
+    if distribute_lib.in_cross_replica_context():
+      raise ValueError('apply_gradients() must be called in a replica context.')
+
+    if not self._doing_dynamic_loss_scaling():
+      return self._optimizer.apply_gradients(grads_and_vars, global_step, name)
+
+    replica_context = distribute_lib.get_replica_context()
+    grads_and_vars = tuple(grads_and_vars)
+
+    # TODO(nluehr) cleanup GraphKeys.TRAIN_OP
+    return replica_context.merge_call(
+        self._distributed_apply, args=(grads_and_vars, global_step, name))
+
+  def _distributed_apply(self,
+                         distribution,
+                         grads_and_vars,
+                         global_step=None,
+                         name=None):
+    """A version of `apply_gradients` for cross replica context.
+
+    When users are in a cross replica strategy, they must call this rather than
+    `apply_gradients()`.
+
+    Args:
+      distribution: a `DistributionStrategy` object.
+      grads_and_vars: List of (gradient, variable) pairs as returned by
+        `compute_gradients()` and then aggregated across replicas.
+      global_step: Optional (mirrored) `Variable` to increment by one after the
+        variables have been updated.
+      name: Optional name for the returned operation. Default to the name passed
+        to the `Optimizer` constructor.
+
+    Returns:
+      An `Operation` that applies the specified gradients across all
+      replicas. If `global_step` was not None, that operation also
+      increments `global_step`
+    """
+    name = name if name is not None else self.get_name()
+    grads = [g for g, _ in grads_and_vars]
+    loss_scale_update_op, should_apply_grads = (self._loss_scale.update(grads))
+
+    def apply_fn():
+      return self._apply_gradients(distribution, grads_and_vars, global_step,
+                                   name + '-wrapped')
+
+    maybe_apply_op = smart_cond.smart_cond(should_apply_grads, apply_fn,
+                                           control_flow_ops.no_op)
+    return control_flow_ops.group(
+        maybe_apply_op, loss_scale_update_op, name=name)
+
+  def _apply_gradients(self, distribution, grads_and_vars, global_step, name):
+    """Unconditionally apply gradients in cross replica context."""
+    update_ops = distribution.extended.call_for_each_replica(
+        self._optimizer.apply_gradients,
+        args=(grads_and_vars, global_step, name))
+    return distribution.group(update_ops)
+
+  def _apply_sparse(self, grad, var):
+    """This function should never be called."""
+    raise RuntimeError('This function should never be called')
+
+  def _apply_dense(self, grad, var):
+    """This function should never be called."""
+    raise RuntimeError('This function should never be called')
+
+  def _resource_apply_sparse(self, grad, handle, indices):
+    """This function should never be called."""
+    raise RuntimeError('This function should never be called')
+
+  def _resource_apply_dense(self, grad, handle):
+    """This function should never be called."""
+    raise RuntimeError('This function should never be called')
+
+  def variables(self):
+    """Returns the variables of the Optimizer."""
+    return (self._optimizer.variables() +
+            list(self._loss_scale._weights.values()))  # pylint: disable=protected-access

tensorflow/python/training/experimental/loss_scale_optimizer_test.py

@@ -0,0 +1,315 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Tests for MixedPrecisionLossScaleOptimizer."""
+
+import os
+
+from absl.testing import parameterized
+from tensorflow.python.checkpoint import checkpoint as trackable_utils
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import mirrored_strategy
+from tensorflow.python.eager import context
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_conversion
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import custom_gradient
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import gradient_descent
+from tensorflow.python.training import momentum
+from tensorflow.python.training.experimental import loss_scale as loss_scale_module
+from tensorflow.python.training.experimental import loss_scale_optimizer
+
+# If called outside any strategy.scope() calls, this will return the default
+# strategy.
+default_strategy_fn = distribute_lib.get_strategy
+
+
+def create_mirrored_strategy():
+  if context.num_gpus() >= 1:
+    return mirrored_strategy.MirroredStrategy(['cpu:0', 'gpu:0'])
+  else:
+    return mirrored_strategy.MirroredStrategy(['cpu:0'])
+
+
+TESTCASES = ({
+    'testcase_name': 'Base',
+    'strategy_fn': default_strategy_fn
+}, {
+    'testcase_name': 'Distribute',
+    'strategy_fn': create_mirrored_strategy
+})
+
+
+def get_gradients(opt, loss, params):
+  grads_and_vars = opt.compute_gradients(loss, params)
+  grads, _ = zip(*grads_and_vars)
+  return grads
+
+
+def create_identity_with_grad_check_fn(expected_gradient, expected_dtype=None):
+  """Returns a function that asserts it's gradient has a certain value.
+
+  This serves as a hook to assert intermediate gradients have a certain value.
+  This returns an identity function. The identity's gradient function is also
+  the identity function, except it asserts that the gradient equals
+  `expected_gradient` and has dtype `expected_dtype`.
+
+  Args:
+    expected_gradient: The gradient function asserts that the gradient is this
+      value.
+    expected_dtype: The gradient function asserts the gradient has this dtype.
+
+  Returns:
+    An identity function whose gradient function asserts the gradient has a
+    certain value.
+  """
+  @custom_gradient.custom_gradient
+  def _identity_with_grad_check(x):
+    """Function that asserts it's gradient has a certain value."""
+    x = array_ops.identity(x)
+    def grad(dx):
+      """Gradient function that asserts the gradient has a certain value."""
+      if expected_dtype:
+        assert dx.dtype == expected_dtype, (
+            'dx.dtype should be %s but is: %s' % (expected_dtype, dx.dtype))
+      expected_tensor = tensor_conversion.convert_to_tensor_v2(
+          expected_gradient, dtype=dx.dtype, name='expected_gradient'
+      )
+      # Control dependency is to ensure input is available. It's possible the
+      # dataset will throw a StopIteration to indicate there is no more data, in
+      # which case we don't want to run the assertion.
+      with ops.control_dependencies([x]):
+        assert_op = check_ops.assert_equal(dx, expected_tensor)
+      with ops.control_dependencies([assert_op]):
+        dx = array_ops.identity(dx)
+      return dx
+    return x, grad
+  # Keras sometimes has trouble serializing Lambda layers with a decorated
+  # function. So we define and return a non-decorated function.
+  def identity_with_grad_check(x):
+    return _identity_with_grad_check(x)
+  return identity_with_grad_check
+
+
+class MixedPrecisionLossScaleOptimizerTest(test.TestCase,
+                                           parameterized.TestCase):
+
+  def _run_if_in_graph_mode(self, val):
+    # Running only in graph mode is useful, because optimizers sometimes return
+    # a value that, in Graph mode, is runnable with self.evaluate. But in Eager
+    # mode, the optimizer already does the computations and the return value
+    # cannot be run.
+    if not context.executing_eagerly():
+      self.evaluate(val)
+
+  def _run_fn_with_grad_check(self, strategy, var, opt, expected_grad):
+    grad_check_fn = create_identity_with_grad_check_fn(
+        expected_grad)
+    loss = lambda: grad_check_fn(var) / strategy.num_replicas_in_sync
+    return lambda: opt.minimize(loss, var_list=[var])
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def testFixedLossScaleAppliedToLossWithMinimize(self, strategy_fn):
+    with strategy_fn().scope() as strategy:
+      var = variables.Variable([5.0])
+      opt = gradient_descent.GradientDescentOptimizer(2.0)
+      loss_scale = 10.
+      opt = loss_scale_optimizer.MixedPrecisionLossScaleOptimizer(
+          opt, loss_scale)
+      # We need num_replicas_in_sync to divide loss_scale, otherwise loss_scale
+      # / strategy.num_replicas_in_sync will not be exact, which could lead to
+      # assertion failures due to rounding issues.
+      self.assertEqual(loss_scale % strategy.num_replicas_in_sync, 0)
+      run_fn = self._run_fn_with_grad_check(
+          strategy, var, opt, loss_scale / strategy.num_replicas_in_sync)
+      run_op = strategy.experimental_run(run_fn)
+      self.evaluate(variables.global_variables_initializer())
+      self._run_if_in_graph_mode(run_op)
+      # The loss is the identity of the variable. Therefore the gradient is 1,
+      # and so the variable will be init_val - grad * lr == 5 - 1 * 2 == 3
+      self.assertAllClose([3.], self.evaluate(var))
+
+  @test_util.deprecated_graph_mode_only
+  def testFixedLossScaleAppliedToLossWithGetGradients(self):
+    var = variables.Variable([2.0])
+    opt = gradient_descent.GradientDescentOptimizer(1.0)
+    loss_scale = 10.
+    opt = loss_scale_optimizer.MixedPrecisionLossScaleOptimizer(opt, loss_scale)
+    grad_check_fn = create_identity_with_grad_check_fn(loss_scale)
+    loss = grad_check_fn(var)
+    run_op = get_gradients(opt, loss, [var])
+    self.evaluate(variables.global_variables_initializer())
+    # This will cause an assertion to run, as
+    # create_identity_with_grad_check_fn added an assertion op.
+    self.evaluate(run_op)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def testDynamicLossScale(self, strategy_fn):
+    strategy = strategy_fn()
+    learning_rate = 2.
+    expected_gradient = resource_variable_ops.ResourceVariable(
+        learning_rate / strategy.num_replicas_in_sync)
+    with strategy.scope():
+      var = variables.Variable([5.0])
+      opt = gradient_descent.GradientDescentOptimizer(learning_rate)
+      loss_scale = loss_scale_module.DynamicLossScale(
+          initial_loss_scale=2, increment_period=1, multiplier=2)
+      opt = loss_scale_optimizer.MixedPrecisionLossScaleOptimizer(
+          opt, loss_scale)
+      self.assertEqual(
+          loss_scale.initial_loss_scale % strategy.num_replicas_in_sync, 0)
+
+      run_fn = self._run_fn_with_grad_check(strategy, var, opt,
+                                            expected_gradient)
+      run_op = strategy.experimental_run(run_fn)
+      self.evaluate(variables.global_variables_initializer())
+      self._run_if_in_graph_mode(run_op)
+      # The loss is the identity of the variable. Therefore the gradient is 1,
+      # and so the variable will be init_val - grad * lr == 5 - 1 * 2 == 3
+      self.assertAllClose([3.], self.evaluate(var))
+
+      # Loss scale will be double, so the expected gradient is also doubled.
+      self.evaluate(
+          expected_gradient.assign(2 * learning_rate /
+                                   strategy.num_replicas_in_sync))
+      run_op = strategy.experimental_run(run_fn)
+      self._run_if_in_graph_mode(run_op)
+      # As before, the 2 is subtracted from the variable, making it's new value
+      # 1.
+      self.assertAllClose([1.], self.evaluate(var))
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def testDynamicUpdate(self, strategy_fn):
+    with strategy_fn().scope() as strategy:
+      var = variables.Variable([1.0, 2.0])
+      opt = gradient_descent.GradientDescentOptimizer(1.0)
+      loss_scale = loss_scale_module.DynamicLossScale(
+          initial_loss_scale=2, increment_period=1, multiplier=2)
+      opt = loss_scale_optimizer.MixedPrecisionLossScaleOptimizer(
+          opt, loss_scale)
+
+      # Test optimizer with finite gradients
+      loss = lambda: var * 2.0 / strategy.num_replicas_in_sync
+      run_fn = lambda: opt.minimize(loss, var_list=[var])
+      run_op = strategy.experimental_run(run_fn)
+      self.evaluate(variables.global_variables_initializer())
+      self._run_if_in_graph_mode(run_op)
+      # Gradient is 2, so variable will have 2 subtracted from it
+      self.assertAllClose([-1.0, 0.0], self.evaluate(var))
+      # Loss scale has doubled from 2 to 4
+      self.assertEqual(4., self.evaluate(opt._loss_scale()))
+
+      # Test optimizer with NaN gradients
+      loss = lambda: var * float('NaN')
+      run_fn = lambda: opt.minimize(loss, var_list=[var])
+      run_op = strategy.experimental_run(run_fn)
+      self._run_if_in_graph_mode(run_op)
+      # Variable should not change from before, due to NaN gradients.
+      self.assertAllClose(self.evaluate(var), [-1.0, 0.0])
+      # Loss scale should half due to NaN gradients.
+      self.assertEqual(2., self.evaluate(opt._loss_scale()))
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def testDynamicLossScaleWithSlots(self, strategy_fn):
+    with strategy_fn().scope() as strategy:
+      var = variables.Variable([1.0, 2.0])
+      # An SGD optimizer with momentum has slot variables.
+      opt = momentum.MomentumOptimizer(1.0, momentum=1.)
+      initial_loss_scale = 2.
+      loss_scale = loss_scale_module.DynamicLossScale(
+          initial_loss_scale=initial_loss_scale,
+          increment_period=1,
+          multiplier=4)
+      opt = loss_scale_optimizer.MixedPrecisionLossScaleOptimizer(
+          opt, loss_scale)
+      loss = lambda: var / strategy.num_replicas_in_sync
+      run_fn = lambda: opt.minimize(loss, var_list=[var])
+      run_op = strategy.experimental_run(run_fn)
+      self.evaluate(variables.global_variables_initializer())
+      self._run_if_in_graph_mode(run_op)
+      # The momentum accumulator starts at 0 and the gradient is 1. The
+      # accumulator is incremented by the gradient, so it is now 1. Then the
+      # variable is subtracted by the accumulator, so the variable is subtracted
+      # by 1.
+      self.assertAllClose([0.0, 1.0], self.evaluate(var))
+      self.assertEqual(self.evaluate(opt._loss_scale()), initial_loss_scale * 4)
+
+      run_op = strategy.experimental_run(run_fn)
+      self._run_if_in_graph_mode(run_op)
+      # The momentum accumulator was 1 before this step and the gradient is 1.
+      # The accumulator is incremented by the gradient, so it is now 2. Then the
+      # variable is subtracted by the accumulator, so the variable is subtracted
+      # by 2.
+      self.assertAllClose([-2., -1.], self.evaluate(var))
+      self.assertEqual(
+          self.evaluate(opt._loss_scale()), initial_loss_scale * 16)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def testCheckpoint(self, strategy_fn):
+    strategy = strategy_fn()
+    if (isinstance(strategy, mirrored_strategy.MirroredStrategy) and
+        not context.executing_eagerly()):
+      # TODO(b/121381184): Enable running the test in this case.
+      return
+
+    with self.test_session(), strategy.scope():
+      # Build and run a simple model.
+      var = variables.Variable([2.0])
+      loss_scale = loss_scale_module.DynamicLossScale(
+          initial_loss_scale=1., increment_period=2., multiplier=2.)
+      opt = momentum.MomentumOptimizer(1.0, momentum=1.)
+      opt = loss_scale_optimizer.MixedPrecisionLossScaleOptimizer(
+          opt, loss_scale)
+      run_fn = lambda: opt.minimize(lambda: var + 1., var_list=[var])
+      opt_op = strategy.experimental_run(run_fn)
+      self.evaluate(variables.global_variables_initializer())
+      self.evaluate(opt_op)
+      self.assertEqual(self.evaluate(loss_scale()), 1.)
+      self.assertEqual(self.evaluate(loss_scale._num_good_steps), 1)
+
+      # Save a checkpoint.
+      checkpoint = trackable_utils.Checkpoint(optimizer=opt)
+      prefix = os.path.join(self.get_temp_dir(), 'ckpt')
+      save_path = checkpoint.save(prefix)
+
+      # Run model again.
+      self.evaluate(strategy.experimental_run(run_fn))
+      self.assertEqual(self.evaluate(loss_scale()), 2.)
+      self.assertEqual(self.evaluate(loss_scale._num_good_steps), 0)
+
+      # Load checkpoint and ensure loss scale is back to it's original value.
+      status = checkpoint.restore(save_path)
+      status.assert_consumed()
+      status.run_restore_ops()
+      self.assertEqual(self.evaluate(loss_scale()), 1.)
+      self.assertEqual(self.evaluate(loss_scale._num_good_steps), 1)
+
+  def testPassingNoneToLossScale(self):
+    opt = gradient_descent.GradientDescentOptimizer(1.0)
+    with self.assertRaisesRegex(ValueError, r'loss_scale cannot be None'):
+      loss_scale_optimizer.MixedPrecisionLossScaleOptimizer(opt, None)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/training/experimental/loss_scale_test.py

@@ -0,0 +1,325 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Tests for LossScale classes.."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import mirrored_strategy
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor as tensor_lib
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training.experimental import loss_scale as loss_scale_module
+
+# TODO(reedwm): Create test case using multiple graphs
+
+# If called outside any strategy.scope() calls, this will return the default
+# strategy.
+default_strategy_fn = distribute_lib.get_strategy
+
+
+def create_mirrored_strategy():
+  if context.num_gpus() >= 1:
+    return mirrored_strategy.MirroredStrategy(['cpu:0', 'gpu:0'])
+  else:
+    return mirrored_strategy.MirroredStrategy(['cpu:0'])
+
+
+TESTCASES = ({
+    'testcase_name': 'base',
+    'strategy_fn': default_strategy_fn
+}, {
+    'testcase_name': 'distribute',
+    'strategy_fn': create_mirrored_strategy
+})
+
+
+class FixedLossScaleTest(test.TestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_basic(self):
+    loss_scale_value = 1000
+    loss_scale = loss_scale_module.FixedLossScale(loss_scale_value)
+
+    update_op, should_apply = loss_scale.update([constant_op.constant(0.)])
+    self.evaluate(update_op)
+    # should_apply should be a bool instead of a tensor, so that a tf.cond does
+    # not have to be built in the graph by the caller.
+    self.assertIsInstance(should_apply, bool)
+    self.assertTrue(should_apply)
+    self.assertEqual(loss_scale_value, self.evaluate(loss_scale()))
+
+    update_op, should_apply = loss_scale.update(
+        [constant_op.constant(float('NaN'))])
+    self.evaluate(update_op)
+    self.assertIsInstance(should_apply, bool)
+    self.assertTrue(should_apply)
+    self.assertEqual(loss_scale_value, self.evaluate(loss_scale()))
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_serialization(self):
+    loss_scale = loss_scale_module.get(123)
+    config = loss_scale.get_config()
+    loss_scale = loss_scale_module.FixedLossScale.from_config(config)
+    self.assertEqual(self.evaluate(loss_scale()), 123.)
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_call_type(self):
+    scalar = loss_scale_module.FixedLossScale(123)
+    self.assertIsInstance(scalar(), tensor_lib.Tensor)
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_repr(self):
+    loss_scale = loss_scale_module.FixedLossScale(123)
+    self.assertEqual(repr(loss_scale), 'FixedLossScale(123.0)')
+
+
+def _get_example_iter(inputs):
+  dataset = dataset_ops.Dataset.from_tensor_slices(inputs)
+  return dataset_ops.make_one_shot_iterator(dataset)
+
+
+class DynamicLossScaleTest(test.TestCase, parameterized.TestCase):
+
+  def _get_tensor(self, is_finite):
+    tensor = cond.cond(is_finite, lambda: 1., lambda: float('NaN'))
+
+    if not distribute_lib.has_strategy():
+      return tensor
+
+    def get():
+      rep_id = (
+          distribute_lib.get_replica_context()
+          .replica_id_in_sync_group)
+      return cond.cond(
+          math_ops.equal(rep_id, 0), lambda: tensor, lambda: 1.)
+
+    distribution = distribute_lib.get_strategy()
+    return distribution.extended.call_for_each_replica(get)
+
+  def _test_helper(self,
+                   inputs,
+                   expected_outputs,
+                   initial_loss_scale=1.,
+                   increment_period=2,
+                   multiplier=2):
+    loss_scale = loss_scale_module.DynamicLossScale(
+        initial_loss_scale=initial_loss_scale,
+        increment_period=increment_period,
+        multiplier=multiplier)
+    itr = _get_example_iter(inputs)
+
+    def update():
+      is_finite = itr.get_next()
+      grad = self._get_tensor(is_finite)
+      update_op, should_apply_gradients = loss_scale.update([grad])
+      assert_op = check_ops.assert_equal(should_apply_gradients, is_finite)
+      if context.executing_eagerly():
+        return
+      with ops.control_dependencies([assert_op]):
+        return array_ops.identity(update_op)
+
+    actual_outputs = []
+
+    if not context.executing_eagerly():
+      update_op = update()
+      self.evaluate(variables.global_variables_initializer())
+    for _ in range(len(inputs)):
+      if context.executing_eagerly():
+        update()
+      else:
+        self.evaluate(update_op)
+      actual_outputs.append(self.evaluate(loss_scale()))
+    self.assertEqual(actual_outputs, expected_outputs)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def test_increase(self, strategy_fn):
+    with strategy_fn().scope():
+      inputs = [True] * 6
+      expected_outputs = [1, 2, 2, 4, 4, 8]
+      self._test_helper(inputs, expected_outputs)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def test_keep_increasing_until_capped(self, strategy_fn):
+    with strategy_fn().scope():
+      init_loss_scale = np.finfo(np.float32).max / 4
+      max_float = np.finfo(np.float32).max
+
+      inputs = [True] * 6
+      # Output is capped the 2nd time it doubles.
+      expected_outputs = [
+          init_loss_scale, init_loss_scale * 2, init_loss_scale * 2, max_float,
+          max_float, max_float
+      ]
+
+      self._test_helper(inputs, expected_outputs, init_loss_scale)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def test_decrease_every_step(self, strategy_fn):
+    with strategy_fn().scope():
+      inputs = [False] * 6
+      init_loss_scale = 1024
+      expected_outputs = [512, 256, 128, 64, 32, 16]
+
+    self._test_helper(inputs, expected_outputs, init_loss_scale)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def test_keep_decreasing_until_one(self, strategy_fn):
+    with strategy_fn().scope():
+      inputs = [False] * 6
+      init_loss_scale = 16
+      expected_outputs = [8, 4, 2, 1, 1, 1]
+
+      self._test_helper(inputs, expected_outputs, init_loss_scale)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def test_nan_clear_good_step(self, strategy_fn):
+    with strategy_fn().scope():
+      inputs = [True, True, True, False, True]
+      expected_outputs = [1, 2, 2, 1, 1]
+      self._test_helper(inputs, expected_outputs)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def test_trigger_loss_scale_update_each_step(self, strategy_fn):
+    with strategy_fn().scope():
+      init_loss_scale = 1
+      increment_period = 1
+
+      inputs = [True] * 3 + [False, True, True]
+      expected_outputs = [2, 4, 8, 4, 8, 16]
+
+      self._test_helper(inputs, expected_outputs, init_loss_scale,
+                        increment_period)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def test_alternating_good_and_bad_gradients_trigger_each_step(
+      self, strategy_fn):
+    with strategy_fn().scope():
+      init_loss_scale = 1
+      increment_period = 1
+
+      inputs = [True, False] * 4 + [True]
+      expected_outputs = [2, 1, 2, 1, 2, 1, 2, 1, 2]
+      self._test_helper(inputs, expected_outputs, init_loss_scale,
+                        increment_period)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def test_alternating_good_and_bad_gradients_trigger_every_other_step(
+      self, strategy_fn):
+    with strategy_fn().scope():
+      init_loss_scale = 32
+      increment_period = 2
+
+      inputs = [True, False] * 3 + [True]
+      expected_outputs = [32, 16, 16, 8, 8, 4, 4]
+      self._test_helper(inputs, expected_outputs, init_loss_scale,
+                        increment_period)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def test_nondefault_multiplier(self, strategy_fn):
+    with strategy_fn().scope():
+      init_loss_scale = 4
+      multiplier = 3
+      inputs = [True, True, False, True, True]
+      expected_outputs = [4, 12, 4, 4, 12]
+      self._test_helper(
+          inputs, expected_outputs, init_loss_scale, multiplier=multiplier)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def test_random_mix_good_and_bad_gradients(self, strategy_fn):
+    with strategy_fn().scope():
+      init_loss_scale = 4
+      inputs = [
+          False, True, True, True, False, True, False, True, True, True, False
+      ]
+      expected_outputs = [2, 2, 4, 4, 2, 2, 1, 1, 2, 2, 1]
+      self._test_helper(inputs, expected_outputs, init_loss_scale)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def test_single_tensor_gradient(self, strategy_fn):
+    with strategy_fn().scope():
+      loss_scale = loss_scale_module.DynamicLossScale()
+      grad = constant_op.constant(4.0)
+      _, should_apply = loss_scale.update(grad)
+      self.assertTrue(self.evaluate(should_apply))
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_serialization(self):
+    loss_scale = loss_scale_module.DynamicLossScale(
+        initial_loss_scale=1, increment_period=2, multiplier=3)
+    config = loss_scale.get_config()
+    loss_scale = loss_scale_module.DynamicLossScale.from_config(config)
+    self.evaluate(variables.global_variables_initializer())
+    self.assertEqual(self.evaluate(loss_scale()), 1)
+    self.assertEqual(loss_scale.increment_period, 2)
+    self.assertEqual(loss_scale.multiplier, 3)
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_update_with_none_gradients(self):
+    loss_scale = loss_scale_module.DynamicLossScale()
+    loss_scale.update([None])
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_get(self):
+    scalar = loss_scale_module.get('dynamic')
+    scalar2 = loss_scale_module.DynamicLossScale()
+    self.assertEqual(scalar.initial_loss_scale, scalar2.initial_loss_scale)
+    self.assertEqual(scalar.increment_period, scalar2.increment_period)
+    self.assertEqual(scalar.multiplier, scalar2.multiplier)
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_call_type(self):
+    scalar = loss_scale_module.DynamicLossScale()
+    self.assertIsInstance(scalar(), tensor_lib.Tensor)
+
+  @parameterized.named_parameters(*TESTCASES)
+  @test_util.run_in_graph_and_eager_modes
+  def test_repr(self, strategy_fn):
+    with strategy_fn().scope():
+      loss_scale = loss_scale_module.DynamicLossScale(
+          initial_loss_scale=1, increment_period=2, multiplier=3)
+      if context.executing_eagerly():
+        self.assertEqual(repr(loss_scale),
+                         'DynamicLossScale(current_loss_scale=1.0, '
+                         'num_good_steps=0, initial_loss_scale=1.0, '
+                         'increment_period=2, multiplier=3.0)')
+      else:
+        self.assertEqual(repr(loss_scale),
+                         'DynamicLossScale(initial_loss_scale=1.0, '
+                         'increment_period=2, multiplier=3.0)')
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/training/experimental/mixed_precision.py

@@ -0,0 +1,248 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Contains functions to use mixed precision with the graph rewrite."""
+
+from tensorflow.python.framework import config
+from tensorflow.python.platform import tf_logging
+from tensorflow.python.training import optimizer
+from tensorflow.python.training.experimental import loss_scale_optimizer as loss_scale_optimizer_v1
+from tensorflow.python.training.experimental import mixed_precision_global_state
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+# A mapping between optimizers and (wrapper_fn, wrapper_cls) pairs. wrapper_cls
+# is a loss scale optimizer class, and wrapper_fn is a function that takes in
+# an optimizer and LossScale and returns a wrapper_cls instance.
+_REGISTERED_WRAPPER_OPTIMIZER_CLS = {
+    optimizer.Optimizer:
+        (loss_scale_optimizer_v1.MixedPrecisionLossScaleOptimizer,) * 2,
+}
+
+
+@tf_export('__internal__.mixed_precision.register_loss_scale_wrapper', v1=[])
+def register_loss_scale_wrapper(optimizer_cls, wrapper_fn, wrapper_cls=None):
+  """Registers a loss scale optimizer wrapper.
+
+  `tf.compat.v1.mixed_precision.enable_mixed_precision_graph_rewrite`
+  automatically wraps an optimizer with an optimizer wrapper that performs loss
+  scaling. This function registers a
+  `(base_cls, wrapper_fn, wrapper_cls)` triple
+  that is used by `enable_mixed_precision_graph_rewrite`, where
+  `wrapper_fn` is called to create a `wrapper_cls` instance that wraps an
+  `optimizer_cls` instance.
+
+  Args:
+    optimizer_cls: A base optimizer class, e.g. `tf.keras.optimizers.Optimizer`.
+    wrapper_fn: A function that takes in arguments "optimizer" and
+      "loss_scale", and returns a loss scale optimizer of type "wrapper_cls"
+      that wraps "optimizer".
+    wrapper_cls: A loss scale optimizer class. Defaults to `wrapper_fn`, in
+      which case `wrapper_fn` should be a loss scale optimizer class whose
+      constructor takes in arguments "optimizer" and "loss_scale".
+  """
+  _REGISTERED_WRAPPER_OPTIMIZER_CLS[optimizer_cls] = (
+      wrapper_fn, wrapper_cls or wrapper_fn)
+
+
+def _wrap_optimizer(opt, loss_scale):
+  """Wraps an optimizer with a LossScaleOptimizer."""
+
+  for _, wrapper_optimizer in _REGISTERED_WRAPPER_OPTIMIZER_CLS.values():
+    if isinstance(opt, wrapper_optimizer):
+      raise ValueError('"opt" must not already be an instance of a {cls}. '
+                       '`enable_mixed_precision_graph_rewrite` will '
+                       'automatically wrap the optimizer with a '
+                       '{cls}.'
+                       .format(cls=wrapper_optimizer.__name__))
+
+  for optimizer_cls, (wrapper_fn, _) in (
+      _REGISTERED_WRAPPER_OPTIMIZER_CLS.items()):
+    if isinstance(opt, optimizer_cls):
+      return wrapper_fn(opt, loss_scale)
+
+  raise ValueError('"opt" must be an instance of a tf.train.Optimizer or a '
+                   'tf.keras.optimizers.Optimizer, but got: %s' % opt)
+
+
+@deprecation.deprecated_endpoints(
+    'train.experimental.enable_mixed_precision_graph_rewrite')
+@tf_export(v1=['mixed_precision.enable_mixed_precision_graph_rewrite',
+               'train.experimental.enable_mixed_precision_graph_rewrite'])
+def enable_mixed_precision_graph_rewrite_v1(opt, loss_scale='dynamic'):
+  """Enable mixed precision via a graph rewrite.
+
+  Mixed precision is the use of both float32 and float16 data types when
+  training a model to improve performance. This is achieved via a graph rewrite
+  operation and a loss-scale optimizer.
+
+  Performing arithmetic operations in float16 takes advantage of specialized
+  processing units, such as NVIDIA Tensor Cores, for much higher arithmetic
+  throughput. However, due to the smaller representable range, performing the
+  entire training with float16 can result in gradient underflow, that is, small
+  gradient values becoming zeroes. Instead, performing only select arithmetic
+  operations in float16 results in higher throughput and decreased training
+  time when using compatible hardware accelerators while also reducing memory
+  usage, typically without sacrificing model accuracy.
+
+  Note: While the mixed precision rewrite changes the datatype of various
+  layers throughout the model, the same accuracy reached in float32 is
+  expected. If a `NaN` gradient occurs with dynamic loss scaling, the model
+  update for that batch is skipped. In this case, the global step count is not
+  incremented, and the `LossScaleOptimizer` attempts to decrease the loss
+  scaling value to avoid `NaN` values in subsequent iterations. This approach
+  has been shown to achieve the same accuracy as float32 and, in most cases,
+  better training throughput.
+
+  Example:
+
+  ```python
+  model = tf.keras.models.Sequential([
+      tf.keras.layers.Dense(64, activation='relu'),
+      tf.keras.layers.Dense(64, activation='softmax'),
+  ])
+
+  opt = tf.keras.optimizers.SGD()
+  opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(opt)
+  model.compile(loss="mse", optimizer=opt)
+
+  x_train = np.random.random((1024, 64))
+  y_train = np.random.random((1024, 64))
+  model.fit(x_train, y_train)
+  ```
+
+  Calling `enable_mixed_precision_graph_rewrite(opt)` enables the graph rewrite
+  operation before computing gradients. The function additionally returns an
+  `Optimizer` (`opt`) wrapped with a `LossScaleOptimizer`. This prevents
+  underflow in the float16 tensors during the backward pass. An optimizer of
+  type `tf.train.Optimizer` or `tf.keras.optimizers.Optimizer` must be passed
+  to this function, which will then be wrapped to use loss scaling.
+
+  The graph rewrite operation changes the `dtype` of certain operations in the
+  graph from float32 to float16. There are several categories of operations
+  that are either included or excluded by this rewrite operation. The following
+  categories of Ops are defined inside corresponding functions under the class
+  `AutoMixedPrecisionLists` in
+  <a href="https://github.com/tensorflow/tensorflow/blob/master/tensorflow/
+  core/grappler/optimizers/auto_mixed_precision_lists.h">
+  auto_mixed_precision_lists.h</a>:
+
+  * `ClearList`: Ops that do not have numerically significant adverse effects.
+  E.g. `ArgMax` and `Floor`.
+  * `AllowList`: Ops that are considered numerically safe for execution in
+  float16, and thus are always converted. E.g. `Conv2D`.
+  * `DenyList`: Ops that are numerically unsafe to execute in float16 and
+  can negatively affect downstream nodes. E.g. `Softmax`.
+  * `GrayList`: Ops that are considered numerically safe for execution in
+  float16 unless downstream from a DenyList Op. E.g. `Add` and `AvgPool`.
+
+  When this function is used, gradients should only be computed and applied
+  with the returned optimizer, either by calling `opt.minimize()` or
+  `opt.compute_gradients()` followed by `opt.apply_gradients()`.
+  Gradients should not be computed with `tf.gradients` or `tf.GradientTape`.
+  This is because the returned optimizer will apply loss scaling, and
+  `tf.gradients` or `tf.GradientTape` will not. If you do directly use
+  `tf.gradients` or `tf.GradientTape`, your model may not converge due to
+  float16 underflow problems.
+
+  When eager execution is enabled, the mixed precision graph rewrite is only
+  enabled within `tf.function`s, as outside `tf.function`s, there is no graph.
+
+  For NVIDIA GPUs with Tensor cores, as a general performance guide, dimensions
+  (such as batch size, input size, output size, and channel counts)
+  should be powers of two if under 256, or  otherwise divisible by 8 if above
+  256. For more information, check out the
+  [NVIDIA Deep Learning Performance Guide](
+  https://docs.nvidia.com/deeplearning/sdk/dl-performance-guide/index.html).
+
+  Currently, mixed precision is only enabled on NVIDIA Tensor Core GPUs with
+  Compute Capability 7.0 and above (Volta, Turing, or newer architectures). The
+  parts of the graph on CPUs and TPUs are untouched by the graph rewrite.
+
+  Raises:
+    `ValueError`, if the `tf.keras.mixed_precision` API is also used by calling
+    `tf.keras.mixed_precision.set_global_policy`. Only one mixed precision
+    API can be used.
+
+  Args:
+    opt: An instance of a `tf.keras.optimizers.Optimizer` or a
+      `tf.train.Optimizer`.
+    loss_scale: Either an int/float, the string `"dynamic"`, or an instance of
+      a `tf.mixed_precision.experimental.LossScale`. The loss scale to use. It
+      is recommended to keep this as its default value of `"dynamic"`, which
+      will adjust the scaling automatically to prevent `Inf` or `NaN` values.
+
+  Returns:
+    A version of `opt` that will use loss scaling to prevent underflow.
+  """
+  if mixed_precision_global_state.is_using_mixed_precision_policy():
+    raise ValueError(
+        'The mixed precision graph rewrite cannot be enabled, because the '
+        'global Keras dtype Policy has been set to a mixed precision policy. '
+        'At most, one of the following can be called:\n\n'
+        '  1. tf.keras.mixed_precision.set_global_policy() with a mixed '
+        'precision policy (You called this first)\n\n'
+        '  2. tf.train.experimental.enable_mixed_precision_graph_rewrite() '
+        '(You called this second)\n'
+        'You called both functions, which is an error, because both functions '
+        'enable you to use mixed precision. If in doubt which function to use, '
+        'use the first, as it supports Eager execution and is more '
+        'customizable.')
+
+  if mixed_precision_global_state.non_mixed_precision_session_created():
+    # TODO(reedwm): Give the stacktrace of the existing Sessions. And if the
+    # Sessions have already been closed, do not raise this error message.
+    tf_logging.warn('You already have existing Sessions that do not use mixed '
+                    'precision. enable_mixed_precision_graph_rewrite() will '
+                    'not affect these Sessions.')
+  opt = _wrap_optimizer(opt, loss_scale)
+  config.set_optimizer_experimental_options({'auto_mixed_precision': True})
+  mixed_precision_global_state.set_mixed_precision_graph_rewrite_enabled(True)
+  return opt
+
+
+@deprecation.deprecated_endpoints(
+    'train.experimental.disable_mixed_precision_graph_rewrite')
+@tf_export(v1=['mixed_precision.disable_mixed_precision_graph_rewrite',
+               'train.experimental.disable_mixed_precision_graph_rewrite'])
+def disable_mixed_precision_graph_rewrite_v1():
+  """Disables the mixed precision graph rewrite.
+
+  After this is called, the mixed precision graph rewrite will no longer run for
+  new Sessions, and so float32 operations will no longer be converted to float16
+  in such Sessions. However, any existing Sessions will continue to have the
+  graph rewrite enabled if they were created after
+  `enable_mixed_precision_graph_rewrite` was called but before
+  `disable_mixed_precision_graph_rewrite` was called.
+
+  This does not undo the effects of loss scaling. Any optimizers wrapped with a
+  LossScaleOptimizer will continue to do loss scaling, although this loss
+  scaling will no longer be useful if the optimizer is used in new Sessions, as
+  the graph rewrite no longer converts the graph to use float16.
+
+  This function is useful for unit testing. A unit tests can test using the
+  mixed precision graph rewrite, then disable it so future unit tests continue
+  using float32. If this is done, unit tests should not share a single session,
+  as `enable_mixed_precision_graph_rewrite` and
+  `disable_mixed_precision_graph_rewrite` have no effect on existing sessions.
+  """
+  # We only have a separate V1 version of this function, because the V1
+  # docstring mentions sessions.
+  if (not
+      mixed_precision_global_state.is_mixed_precision_graph_rewrite_enabled()):
+    tf_logging.warn('disable_mixed_precision_graph_rewrite() called when mixed '
+                    'precision is already disabled.')
+  config.set_optimizer_experimental_options({'auto_mixed_precision': False})
+  mixed_precision_global_state.set_mixed_precision_graph_rewrite_enabled(False)

tensorflow/python/training/experimental/mixed_precision_global_state.py

@@ -0,0 +1,66 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Contains global variables related to mixed precision.
+
+This is not part of mixed_precision.py to avoid a circular dependency.
+mixed_precision.py depends on Session, and Session depends on this file.
+"""
+
+from tensorflow.python.util.tf_export import tf_export
+
+# Whether the mixed precision graph rewrite has been enabled or not with
+# `enable_mixed_precision_graph_rewrite`. Used to turn on auto_mixed_precision
+# in ConfigProtos passed to Sessions.
+_mixed_precision_graph_rewrite_is_enabled = False
+
+
+# True if a Session has been created without the mixed precision graph rewrite
+# being enabled. Used to give a warning if mixed precision is enabled after a
+# Session has already been created.
+_non_mixed_precision_session_created = False
+
+# Whether the global tf.keras.mixed_precision.Policy uses mixed precision. Used
+# to raise an error message if both a mixed Policy and the graph rewrite are
+# used at the same time.
+_using_mixed_precision_policy = False
+
+
+@tf_export('__internal__.train.is_mixed_precision_graph_rewrite_enabled', v1=[])
+def is_mixed_precision_graph_rewrite_enabled():
+  return _mixed_precision_graph_rewrite_is_enabled
+
+
+def set_mixed_precision_graph_rewrite_enabled(enabled):
+  global _mixed_precision_graph_rewrite_is_enabled
+  _mixed_precision_graph_rewrite_is_enabled = enabled
+
+
+def non_mixed_precision_session_created():
+  return _non_mixed_precision_session_created
+
+
+def set_non_mixed_precision_session_created(created):
+  global _non_mixed_precision_session_created
+  _non_mixed_precision_session_created = created
+
+
+def is_using_mixed_precision_policy():
+  return _using_mixed_precision_policy
+
+
+@tf_export('__internal__.train.set_using_mixed_precision_policy', v1=[])
+def set_using_mixed_precision_policy(is_using):
+  global _using_mixed_precision_policy
+  _using_mixed_precision_policy = is_using

tensorflow/python/training/experimental/mixed_precision_test.py

@@ -0,0 +1,207 @@
+# Copyright 2019 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.
+# ==============================================================================
+
+import os
+
+from absl.testing import parameterized
+
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.client import session
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import config
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging
+from tensorflow.python.training import gradient_descent as gradient_descent_v1
+from tensorflow.python.training.experimental import loss_scale_optimizer as loss_scale_optimizer_v1
+from tensorflow.python.training.experimental import mixed_precision
+from tensorflow.python.training.experimental import mixed_precision_global_state
+
+
+class MixedPrecisionTest(test.TestCase, parameterized.TestCase):
+
+  IGNORE_PERF_VAR = 'TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_IGNORE_PERFORMANCE'
+
+  def setUp(self):
+    super(MixedPrecisionTest, self).setUp()
+    # Enable the tests to be run on pre-Volta GPUs by telling the grappler pass
+    # to ignore performance and always transform the graph.
+    self._original_ignore_perf_value = os.getenv(self.IGNORE_PERF_VAR)
+    os.environ[self.IGNORE_PERF_VAR] = '1'
+
+  def tearDown(self):
+    # Set the IGNORE_PERF_VAR variable back to it's original value.
+    if self._original_ignore_perf_value is not None:
+      os.environ[self.IGNORE_PERF_VAR] = self._original_ignore_perf_value
+    else:
+      del os.environ[self.IGNORE_PERF_VAR]
+
+    mixed_precision.disable_mixed_precision_graph_rewrite_v1()
+    super(MixedPrecisionTest, self).tearDown()
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_wrap_optimizer(self):
+    opt = gradient_descent_v1.GradientDescentOptimizer(1.0)
+    opt = mixed_precision.enable_mixed_precision_graph_rewrite_v1(opt, 123.)
+    self.assertIsInstance(
+        opt, loss_scale_optimizer_v1.MixedPrecisionLossScaleOptimizer)
+    self.assertEqual(self.evaluate(opt._loss_scale()), 123.)
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_optimizer_errors(self):
+    opt = 1
+    expected_regex = ('"opt" must be an instance of a tf.train.Optimizer or '
+                      'a tf.keras.optimizers.Optimizer, but got')
+    with self.assertRaisesRegex(ValueError, expected_regex):
+      mixed_precision.enable_mixed_precision_graph_rewrite_v1(opt)
+    self.assertFalse(config.get_optimizer_experimental_options()
+                     .get('auto_mixed_precision', False))
+
+    opt = gradient_descent_v1.GradientDescentOptimizer(1.0)
+    opt = loss_scale_optimizer_v1.MixedPrecisionLossScaleOptimizer(opt,
+                                                                   'dynamic')
+    with self.assertRaisesRegex(
+        ValueError, '"opt" must not already be an instance of a '
+        'MixedPrecisionLossScaleOptimizer.'):
+      mixed_precision.enable_mixed_precision_graph_rewrite_v1(opt)
+    self.assertFalse(config.get_optimizer_experimental_options()
+                     .get('auto_mixed_precision', False))
+
+  @test_util.run_in_graph_and_eager_modes()
+  def test_register_loss_scale_wrapper_with_2_arguments(self):
+    class MyOptimizer:
+      pass
+
+    class MyLossScaleOptimizer(MyOptimizer):
+
+      def __init__(self, inner_optimizer, loss_scale):
+        self.inner_optimizer = inner_optimizer
+        self.loss_scale = loss_scale
+
+    mixed_precision.register_loss_scale_wrapper(MyOptimizer,
+                                                MyLossScaleOptimizer)
+    opt = MyOptimizer()
+    opt = mixed_precision.enable_mixed_precision_graph_rewrite_v1(opt, 123.)
+    self.assertIsInstance(opt, MyLossScaleOptimizer)
+    self.assertEqual(opt.loss_scale, 123.)
+
+  @test_util.run_in_graph_and_eager_modes()
+  def test_register_loss_scale_wrapper_with_3_arguments(self):
+    class MyOptimizer:
+      pass
+
+    class MyLossScaleOptimizer(MyOptimizer):
+
+      def __init__(self, inner_optimizer, loss_scale):
+        self.inner_optimizer = inner_optimizer
+        self.loss_scale = loss_scale
+
+    is_called = False
+
+    def create_lso(inner_optimizer, loss_scale):
+      nonlocal is_called
+      is_called = True
+      return MyLossScaleOptimizer(inner_optimizer, loss_scale)
+
+    mixed_precision.register_loss_scale_wrapper(MyOptimizer,
+                                                create_lso,
+                                                MyLossScaleOptimizer)
+    opt = MyOptimizer()
+    opt = mixed_precision.enable_mixed_precision_graph_rewrite_v1(opt, 123.)
+    self.assertIsInstance(opt, MyLossScaleOptimizer)
+    self.assertEqual(opt.loss_scale, 123.)
+    self.assertTrue(is_called)
+
+  @test_util.run_gpu_only
+  @test_util.run_in_graph_and_eager_modes
+  @test_util.disable_tfrt('Grappler rewrite doesn\'t apply to tfrt.')
+  def test_grappler_pass_enabled(self):
+    opt = gradient_descent_v1.GradientDescentOptimizer(1.0)
+    mixed_precision.enable_mixed_precision_graph_rewrite_v1(opt, 123.)
+
+    var = variables.Variable([[1.0]])
+
+    def overflow_in_float16():
+      out = var * 2 ** 10
+      out = math_ops.matmul(out, out)
+      return array_ops.reshape(out, ())
+
+    if context.executing_eagerly():
+      f = def_function.function(overflow_in_float16)
+      self.assertEqual(f().numpy(), float('Inf'))
+      # Outside a def_function.function, the grappler pass will not be applied.
+      self.assertAlmostEqual(overflow_in_float16().numpy(), 2 ** 20)
+
+      # Test disabling mixed precision.
+      mixed_precision.disable_mixed_precision_graph_rewrite_v1()
+      self.assertEqual(f().numpy(), 2 ** 20)
+    else:
+      with session.Session() as sess:
+        out = overflow_in_float16()
+        sess.run(var.initializer)
+        self.assertEqual(sess.run(out), float('Inf'))
+
+      # Test Session will enable the auto_mixed_precision grappler pass in a
+      # ConfigProto passed by the user
+      with session.Session(config=config_pb2.ConfigProto()) as sess:
+        out = overflow_in_float16()
+        sess.run(var.initializer)
+        self.assertEqual(sess.run(out), float('Inf'))
+
+      # Test disabling mixed precision.
+      mixed_precision.disable_mixed_precision_graph_rewrite_v1()
+      with session.Session() as sess:
+        out = overflow_in_float16()
+        sess.run(var.initializer)
+        self.assertAlmostEqual(sess.run(out), 2 ** 20)
+
+  @test.mock.patch.object(tf_logging, 'warn')
+  def test_warn_if_session_already_exists(self, mock_warn):
+    # Set this to False, so Sessions created in previous tests do not trigger
+    # the warning.
+    mixed_precision_global_state.set_non_mixed_precision_session_created(False)
+
+    with session.Session():
+      mixed_precision.enable_mixed_precision_graph_rewrite_v1(
+          gradient_descent_v1.GradientDescentOptimizer(1.0))
+      mock_warn.assert_any_call(
+          'You already have existing Sessions that do not use mixed precision. '
+          'enable_mixed_precision_graph_rewrite() will not affect these '
+          'Sessions.')
+
+  @test.mock.patch.object(tf_logging, 'warn')
+  def test_do_not_warn_if_session_does_not_already_exist(self, mock_warn):
+    # Set this to False, so Sessions created in previous tests do not trigger
+    # the warning.
+    mixed_precision_global_state.set_non_mixed_precision_session_created(False)
+
+    mixed_precision.enable_mixed_precision_graph_rewrite_v1(
+        gradient_descent_v1.GradientDescentOptimizer(1.0))
+    with session.Session():
+      # Make sure the "You already have existing Sessions" warning was not
+      # issued, since the Session was only created after
+      # enable_mixed_precision_graph_rewrite.
+      for call_arg in mock_warn.call_args_list:
+        msg = call_arg[0][0]
+        self.assertNotIn('You already have existing Sessions that do not use '
+                         'mixed precision', msg)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/training/ftrl.py

@@ -0,0 +1,291 @@
+# 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.
+# ==============================================================================
+"""Ftrl-proximal for TensorFlow."""
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_training_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.training import optimizer
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(v1=["train.FtrlOptimizer"])
+class FtrlOptimizer(optimizer.Optimizer):
+  """Optimizer that implements the FTRL algorithm.
+
+  This version has support for both online L2 (McMahan et al., 2013) and
+  shrinkage-type L2, which is the addition of an L2 penalty
+  to the loss function.
+
+  References:
+    Ad-click prediction:
+      [McMahan et al., 2013](https://dl.acm.org/citation.cfm?id=2488200)
+      ([pdf](https://dl.acm.org/ft_gateway.cfm?id=2488200&ftid=1388399&dwn=1&CFID=32233078&CFTOKEN=d60fe57a294c056a-CB75C374-F915-E7A6-1573FBBC7BF7D526))
+  """
+
+  def __init__(self,
+               learning_rate,
+               learning_rate_power=-0.5,
+               initial_accumulator_value=0.1,
+               l1_regularization_strength=0.0,
+               l2_regularization_strength=0.0,
+               use_locking=False,
+               name="Ftrl",
+               accum_name=None,
+               linear_name=None,
+               l2_shrinkage_regularization_strength=0.0,
+               beta=None):
+    r"""Construct a new FTRL optimizer.
+
+    Args:
+      learning_rate: A float value or a constant float `Tensor`.
+      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. See section 3.1 in (McMahan et al., 2013).
+      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.
+      l2_regularization_strength: A float value, must be greater than or
+        equal to zero.
+      use_locking: If `True` use locks for update operations.
+      name: Optional name prefix for the operations created when applying
+        gradients.  Defaults to "Ftrl".
+      accum_name: The suffix for the variable that keeps the gradient squared
+        accumulator.  If not present, defaults to name.
+      linear_name: The suffix for the variable that keeps the linear gradient
+        accumulator.  If not present, defaults to name + "_1".
+      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.
+        The FTRL formulation can be written as:
+        w_{t+1} = argmin_w(\hat{g}_{1:t}w + L1*||w||_1 + L2*||w||_2^2), where
+        \hat{g} = g + (2*L2_shrinkage*w), and g is the gradient of the loss
+        function w.r.t. the weights w.
+        Specifically, in the absence of L1 regularization, it is equivalent to
+        the following update rule:
+        w_{t+1} = w_t - lr_t / (beta + 2*L2*lr_t) * g_t -
+                  2*L2_shrinkage*lr_t / (beta + 2*L2*lr_t) * w_t
+        where lr_t is the learning rate at t.
+        When input is sparse shrinkage will only happen on the active weights.
+      beta: A float value; corresponds to the beta parameter in the paper.
+
+    Raises:
+      ValueError: If one of the arguments is invalid.
+
+    References:
+      Ad-click prediction:
+        [McMahan et al., 2013](https://dl.acm.org/citation.cfm?id=2488200)
+        ([pdf](https://dl.acm.org/ft_gateway.cfm?id=2488200&ftid=1388399&dwn=1&CFID=32233078&CFTOKEN=d60fe57a294c056a-CB75C374-F915-E7A6-1573FBBC7BF7D526))
+    """
+    super(FtrlOptimizer, self).__init__(use_locking, name)
+
+    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._learning_rate = learning_rate
+    self._learning_rate_power = learning_rate_power
+    self._initial_accumulator_value = initial_accumulator_value
+    self._l1_regularization_strength = l1_regularization_strength
+    self._l2_regularization_strength = l2_regularization_strength
+    self._beta = (0.0 if beta is None else beta)
+    self._l2_shrinkage_regularization_strength = (
+        l2_shrinkage_regularization_strength)
+    self._learning_rate_tensor = None
+    self._learning_rate_power_tensor = None
+    self._l1_regularization_strength_tensor = None
+    self._adjusted_l2_regularization_strength_tensor = None
+    self._l2_shrinkage_regularization_strength_tensor = None
+    self._accum_name = accum_name
+    self._linear_name = linear_name
+
+  def _create_slots(self, var_list):
+    # Create the "accum" and "linear" slots.
+    def _accum_initializer(shape, dtype=dtypes.float32, partition_info=None):
+      del partition_info
+      return array_ops.ones(
+          shape=shape, dtype=dtype) * self._initial_accumulator_value
+    for v in var_list:
+      self._get_or_make_slot_with_initializer(
+          v, _accum_initializer, v.shape, v.dtype, "accum",
+          self._accum_name or self._name)
+      self._zeros_slot(v, "linear", self._linear_name or self._name)
+
+  def _prepare(self):
+    self._learning_rate_tensor = ops.convert_to_tensor(
+        self._learning_rate, name="learning_rate")
+    self._l1_regularization_strength_tensor = ops.convert_to_tensor(
+        self._l1_regularization_strength, name="l1_regularization_strength")
+    # L2 regularization strength with beta added in so that the underlying
+    # TensorFlow ops do not need to include that parameter.
+    self._adjusted_l2_regularization_strength_tensor = ops.convert_to_tensor(
+        self._l2_regularization_strength + self._beta /
+        (2. * math_ops.maximum(self._learning_rate, 1e-36)),
+        name="adjusted_l2_regularization_strength")
+    assert self._adjusted_l2_regularization_strength_tensor is not None
+    self._beta_tensor = ops.convert_to_tensor(self._beta, name="beta")
+    self._l2_shrinkage_regularization_strength_tensor = ops.convert_to_tensor(
+        self._l2_shrinkage_regularization_strength,
+        name="l2_shrinkage_regularization_strength")
+    self._learning_rate_power_tensor = ops.convert_to_tensor(
+        self._learning_rate_power, name="learning_rate_power")
+
+  def _apply_dense(self, grad, var):
+    accum = self.get_slot(var, "accum")
+    linear = self.get_slot(var, "linear")
+    if self._l2_shrinkage_regularization_strength <= 0.0:
+      return gen_training_ops.apply_ftrl(
+          var,
+          accum,
+          linear,
+          grad,
+          math_ops.cast(self._learning_rate_tensor, var.dtype.base_dtype),
+          math_ops.cast(self._l1_regularization_strength_tensor,
+                        var.dtype.base_dtype),
+          math_ops.cast(self._adjusted_l2_regularization_strength_tensor,
+                        var.dtype.base_dtype),
+          math_ops.cast(self._learning_rate_power_tensor, var.dtype.base_dtype),
+          use_locking=self._use_locking)
+    else:
+      return gen_training_ops.apply_ftrl_v2(
+          var,
+          accum,
+          linear,
+          grad,
+          math_ops.cast(self._learning_rate_tensor, var.dtype.base_dtype),
+          math_ops.cast(self._l1_regularization_strength_tensor,
+                        var.dtype.base_dtype),
+          math_ops.cast(self._adjusted_l2_regularization_strength_tensor,
+                        var.dtype.base_dtype),
+          math_ops.cast(self._l2_shrinkage_regularization_strength_tensor,
+                        var.dtype.base_dtype),
+          math_ops.cast(self._learning_rate_power_tensor, var.dtype.base_dtype),
+          use_locking=self._use_locking)
+
+  def _resource_apply_dense(self, grad, var):
+    accum = self.get_slot(var, "accum")
+    linear = self.get_slot(var, "linear")
+    if self._l2_shrinkage_regularization_strength <= 0.0:
+      return gen_training_ops.resource_apply_ftrl(
+          var.handle,
+          accum.handle,
+          linear.handle,
+          grad,
+          math_ops.cast(self._learning_rate_tensor, var.dtype.base_dtype),
+          math_ops.cast(self._l1_regularization_strength_tensor,
+                        var.dtype.base_dtype),
+          math_ops.cast(self._adjusted_l2_regularization_strength_tensor,
+                        var.dtype.base_dtype),
+          math_ops.cast(self._learning_rate_power_tensor, var.dtype.base_dtype),
+          use_locking=self._use_locking)
+    else:
+      return gen_training_ops.resource_apply_ftrl_v2(
+          var.handle,
+          accum.handle,
+          linear.handle,
+          grad,
+          math_ops.cast(self._learning_rate_tensor, var.dtype.base_dtype),
+          math_ops.cast(self._l1_regularization_strength_tensor,
+                        var.dtype.base_dtype),
+          math_ops.cast(self._adjusted_l2_regularization_strength_tensor,
+                        var.dtype.base_dtype),
+          math_ops.cast(self._l2_shrinkage_regularization_strength_tensor,
+                        var.dtype.base_dtype),
+          math_ops.cast(self._learning_rate_power_tensor, var.dtype.base_dtype),
+          use_locking=self._use_locking)
+
+  def _apply_sparse(self, grad, var):
+    accum = self.get_slot(var, "accum")
+    linear = self.get_slot(var, "linear")
+    if self._l2_shrinkage_regularization_strength <= 0.0:
+      return gen_training_ops.sparse_apply_ftrl(
+          var,
+          accum,
+          linear,
+          grad.values,
+          grad.indices,
+          math_ops.cast(self._learning_rate_tensor, var.dtype.base_dtype),
+          math_ops.cast(self._l1_regularization_strength_tensor,
+                        var.dtype.base_dtype),
+          math_ops.cast(self._adjusted_l2_regularization_strength_tensor,
+                        var.dtype.base_dtype),
+          math_ops.cast(self._learning_rate_power_tensor, var.dtype.base_dtype),
+          use_locking=self._use_locking)
+    else:
+      return gen_training_ops.sparse_apply_ftrl_v2(
+          var,
+          accum,
+          linear,
+          grad.values,
+          grad.indices,
+          math_ops.cast(self._learning_rate_tensor, var.dtype.base_dtype),
+          math_ops.cast(self._l1_regularization_strength_tensor,
+                        var.dtype.base_dtype),
+          math_ops.cast(self._adjusted_l2_regularization_strength_tensor,
+                        var.dtype.base_dtype),
+          math_ops.cast(self._l2_shrinkage_regularization_strength_tensor,
+                        grad.dtype.base_dtype),
+          math_ops.cast(self._learning_rate_power_tensor, var.dtype.base_dtype),
+          use_locking=self._use_locking)
+
+  def _resource_apply_sparse(self, grad, var, indices):
+    accum = self.get_slot(var, "accum")
+    linear = self.get_slot(var, "linear")
+    if self._l2_shrinkage_regularization_strength <= 0.0:
+      return gen_training_ops.resource_sparse_apply_ftrl(
+          var.handle,
+          accum.handle,
+          linear.handle,
+          grad,
+          indices,
+          math_ops.cast(self._learning_rate_tensor, grad.dtype),
+          math_ops.cast(self._l1_regularization_strength_tensor, grad.dtype),
+          math_ops.cast(self._adjusted_l2_regularization_strength_tensor,
+                        grad.dtype),
+          math_ops.cast(self._learning_rate_power_tensor, grad.dtype),
+          use_locking=self._use_locking)
+    else:
+      return gen_training_ops.resource_sparse_apply_ftrl_v2(
+          var.handle,
+          accum.handle,
+          linear.handle,
+          grad,
+          indices,
+          math_ops.cast(self._learning_rate_tensor, grad.dtype),
+          math_ops.cast(self._l1_regularization_strength_tensor, grad.dtype),
+          math_ops.cast(self._adjusted_l2_regularization_strength_tensor,
+                        grad.dtype),
+          math_ops.cast(self._l2_shrinkage_regularization_strength_tensor,
+                        grad.dtype),
+          math_ops.cast(self._learning_rate_power_tensor, grad.dtype),
+          use_locking=self._use_locking)

tensorflow/python/training/ftrl_test.py

@@ -0,0 +1,510 @@
+# 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 tests for Ftrl operations."""
+
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import embedding_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import adagrad
+from tensorflow.python.training import ftrl
+from tensorflow.python.training import gradient_descent
+
+
+class FtrlOptimizerTest(test.TestCase):
+
+  def doTestFtrlwithoutRegularization(self, use_resource=False):
+    # The v1 optimizers do not support eager execution
+    with ops.Graph().as_default():
+      for dtype in [dtypes.half, dtypes.float32]:
+        with self.cached_session():
+          if use_resource:
+            var0 = resource_variable_ops.ResourceVariable([0.0, 0.0],
+                                                          dtype=dtype)
+            var1 = resource_variable_ops.ResourceVariable([0.0, 0.0],
+                                                          dtype=dtype)
+          else:
+            var0 = variables.Variable([0.0, 0.0], dtype=dtype)
+            var1 = variables.Variable([0.0, 0.0], dtype=dtype)
+          grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+          grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+          opt = ftrl.FtrlOptimizer(
+              3.0,
+              initial_accumulator_value=0.1,
+              l1_regularization_strength=0.0,
+              l2_regularization_strength=0.0)
+          update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+          self.evaluate(variables.global_variables_initializer())
+
+          v0_val, v1_val = self.evaluate([var0, var1])
+          self.assertAllClose([0.0, 0.0], v0_val)
+          self.assertAllClose([0.0, 0.0], v1_val)
+
+          # Run 3 steps FTRL
+          for _ in range(3):
+            update.run()
+
+          v0_val, v1_val = self.evaluate([var0, var1])
+          self.assertAllCloseAccordingToType(
+              np.array([-2.60260963, -4.29698515]), v0_val, half_rtol=1e-2)
+          self.assertAllCloseAccordingToType(
+              np.array([-0.28432083, -0.56694895]), v1_val)
+
+  def testFtrlWithoutRegularization(self):
+    self.doTestFtrlwithoutRegularization(use_resource=False)
+
+  def testResourceFtrlWithoutRegularization(self):
+    self.doTestFtrlwithoutRegularization(use_resource=True)
+
+  def testFtrlwithoutRegularization2(self):
+    # The v1 optimizers do not support eager execution
+    with ops.Graph().as_default():
+      for dtype in [dtypes.half, dtypes.float32]:
+        with self.cached_session():
+          var0 = variables.Variable([1.0, 2.0], dtype=dtype)
+          var1 = variables.Variable([4.0, 3.0], dtype=dtype)
+          grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+          grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+
+          opt = ftrl.FtrlOptimizer(
+              3.0,
+              initial_accumulator_value=0.1,
+              l1_regularization_strength=0.0,
+              l2_regularization_strength=0.0)
+          update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+          self.evaluate(variables.global_variables_initializer())
+
+          v0_val, v1_val = self.evaluate([var0, var1])
+          self.assertAllCloseAccordingToType([1.0, 2.0], v0_val)
+          self.assertAllCloseAccordingToType([4.0, 3.0], v1_val)
+
+          # Run 3 steps FTRL
+          for _ in range(3):
+            update.run()
+          v0_val, v1_val = self.evaluate([var0, var1])
+          self.assertAllCloseAccordingToType(
+              np.array([-2.55607247, -3.98729396]), v0_val)
+          self.assertAllCloseAccordingToType(
+              np.array([-0.28232238, -0.56096673]), v1_val)
+
+  def testMinimizeSparseResourceVariable(self):
+    # The v1 optimizers do not support eager execution
+    with ops.Graph().as_default():
+      for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+        with self.cached_session():
+          var0 = resource_variable_ops.ResourceVariable([[1.0, 2.0]],
+                                                        dtype=dtype)
+          x = constant_op.constant([[4.0], [5.0]], dtype=dtype)
+          pred = math_ops.matmul(embedding_ops.embedding_lookup([var0], [0]), x)
+          loss = pred * pred
+          sgd_op = ftrl.FtrlOptimizer(1.0).minimize(loss)
+          self.evaluate(variables.global_variables_initializer())
+          # Fetch params to validate initial values
+          self.assertAllCloseAccordingToType([[1.0, 2.0]], self.evaluate(var0))
+          # Run 1 step of sgd
+          sgd_op.run()
+          # Validate updated params
+          self.assertAllCloseAccordingToType([[0, 1]],
+                                             self.evaluate(var0),
+                                             atol=0.01)
+
+  def testFtrlWithL1(self):
+    # The v1 optimizers do not support eager execution
+    with ops.Graph().as_default():
+      for dtype in [dtypes.half, dtypes.float32]:
+        with self.cached_session():
+          var0 = variables.Variable([1.0, 2.0], dtype=dtype)
+          var1 = variables.Variable([4.0, 3.0], dtype=dtype)
+          grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+          grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+
+          opt = ftrl.FtrlOptimizer(
+              3.0,
+              initial_accumulator_value=0.1,
+              l1_regularization_strength=0.001,
+              l2_regularization_strength=0.0)
+          update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+          self.evaluate(variables.global_variables_initializer())
+
+          v0_val, v1_val = self.evaluate([var0, var1])
+          self.assertAllCloseAccordingToType([1.0, 2.0], v0_val)
+          self.assertAllCloseAccordingToType([4.0, 3.0], v1_val)
+
+          # Run 10 steps FTRL
+          for _ in range(10):
+            update.run()
+          v0_val, v1_val = self.evaluate([var0, var1])
+          self.assertAllCloseAccordingToType(
+              np.array([-7.66718769, -10.91273689]), v0_val)
+          self.assertAllCloseAccordingToType(
+              np.array([-0.93460727, -1.86147261]), v1_val)
+
+  def testFtrlWithBeta(self):
+    # The v1 optimizers do not support eager execution
+    with ops.Graph().as_default():
+      for dtype in [dtypes.half, dtypes.float32]:
+        with self.cached_session():
+          var0 = variables.Variable([1.0, 2.0], dtype=dtype)
+          var1 = variables.Variable([4.0, 3.0], dtype=dtype)
+          grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+          grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+
+          opt = ftrl.FtrlOptimizer(3.0, initial_accumulator_value=0.1, beta=0.1)
+          update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+          self.evaluate(variables.global_variables_initializer())
+
+          v0_val, v1_val = self.evaluate([var0, var1])
+          self.assertAllCloseAccordingToType([1.0, 2.0], v0_val)
+          self.assertAllCloseAccordingToType([4.0, 3.0], v1_val)
+
+          # Run 10 steps FTRL
+          for _ in range(10):
+            update.run()
+          v0_val, v1_val = self.evaluate([var0, var1])
+          self.assertAllCloseAccordingToType(
+              np.array([-6.096838, -9.162214]), v0_val)
+          self.assertAllCloseAccordingToType(
+              np.array([-0.717741, -1.425132]), v1_val)
+
+  def testFtrlWithL2_Beta(self):
+    # The v1 optimizers do not support eager execution
+    with ops.Graph().as_default():
+      for dtype in [dtypes.half, dtypes.float32]:
+        with self.cached_session():
+          var0 = variables.Variable([1.0, 2.0], dtype=dtype)
+          var1 = variables.Variable([4.0, 3.0], dtype=dtype)
+          grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+          grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+
+          opt = ftrl.FtrlOptimizer(
+              3.0,
+              initial_accumulator_value=0.1,
+              l1_regularization_strength=0.0,
+              l2_regularization_strength=0.1,
+              beta=0.1)
+          update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+          self.evaluate(variables.global_variables_initializer())
+
+          v0_val, v1_val = self.evaluate([var0, var1])
+          self.assertAllCloseAccordingToType([1.0, 2.0], v0_val)
+          self.assertAllCloseAccordingToType([4.0, 3.0], v1_val)
+
+          # Run 10 steps FTRL
+          for _ in range(10):
+            update.run()
+          v0_val, v1_val = self.evaluate([var0, var1])
+          self.assertAllCloseAccordingToType(
+              np.array([-2.735487, -4.704625]), v0_val)
+          self.assertAllCloseAccordingToType(
+              np.array([-0.294335, -0.586556]), v1_val)
+
+  def testFtrlWithL1_L2(self):
+    # The v1 optimizers do not support eager execution
+    with ops.Graph().as_default():
+      for dtype in [dtypes.half, dtypes.float32]:
+        with self.cached_session():
+          var0 = variables.Variable([1.0, 2.0], dtype=dtype)
+          var1 = variables.Variable([4.0, 3.0], dtype=dtype)
+          grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+          grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+
+          opt = ftrl.FtrlOptimizer(
+              3.0,
+              initial_accumulator_value=0.1,
+              l1_regularization_strength=0.001,
+              l2_regularization_strength=2.0)
+          update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+          self.evaluate(variables.global_variables_initializer())
+
+          v0_val, v1_val = self.evaluate([var0, var1])
+          self.assertAllCloseAccordingToType([1.0, 2.0], v0_val)
+          self.assertAllCloseAccordingToType([4.0, 3.0], v1_val)
+
+          # Run 10 steps FTRL
+          for _ in range(10):
+            update.run()
+
+          v0_val, v1_val = self.evaluate([var0, var1])
+          self.assertAllCloseAccordingToType(
+              np.array([-0.24059935, -0.46829352]), v0_val)
+          self.assertAllCloseAccordingToType(
+              np.array([-0.02406147, -0.04830509]), v1_val)
+
+  def testFtrlWithL1_L2_L2Shrinkage(self):
+    """Test the new FTRL op with support for l2 shrinkage.
+
+    The addition of this parameter which places a constant pressure on weights
+    towards the origin causes the gradient descent trajectory to differ. The
+    weights will tend to have smaller magnitudes with this parameter set.
+    """
+    # The v1 optimizers do not support eager execution
+    with ops.Graph().as_default():
+      for dtype in [dtypes.half, dtypes.float32]:
+        with self.cached_session():
+          var0 = variables.Variable([1.0, 2.0], dtype=dtype)
+          var1 = variables.Variable([4.0, 3.0], dtype=dtype)
+          grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+          grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+
+          opt = ftrl.FtrlOptimizer(
+              3.0,
+              initial_accumulator_value=0.1,
+              l1_regularization_strength=0.001,
+              l2_regularization_strength=2.0,
+              l2_shrinkage_regularization_strength=0.1)
+          update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+          self.evaluate(variables.global_variables_initializer())
+
+          v0_val, v1_val = self.evaluate([var0, var1])
+          self.assertAllCloseAccordingToType([1.0, 2.0], v0_val)
+          self.assertAllCloseAccordingToType([4.0, 3.0], v1_val)
+
+          # Run 10 steps FTRL
+          for _ in range(10):
+            update.run()
+
+          v0_val, v1_val = self.evaluate([var0, var1])
+          self.assertAllCloseAccordingToType(
+              np.array([-0.22578995, -0.44345796]), v0_val)
+          self.assertAllCloseAccordingToType(
+              np.array([-0.14378493, -0.13229476]), v1_val)
+
+  def testFtrlWithL1_L2_L2ShrinkageSparse(self):
+    """Tests the new FTRL op with support for l2 shrinkage on sparse grads."""
+    # The v1 optimizers do not support eager execution
+    with ops.Graph().as_default():
+      for dtype in [dtypes.half, dtypes.float32]:
+        with self.cached_session():
+          var0 = variables.Variable([[1.0], [2.0]], dtype=dtype)
+          var1 = variables.Variable([[4.0], [3.0]], dtype=dtype)
+          grads0 = indexed_slices.IndexedSlices(
+              constant_op.constant([0.1], shape=[1, 1], dtype=dtype),
+              constant_op.constant([0]), constant_op.constant([2, 1]))
+          grads1 = indexed_slices.IndexedSlices(
+              constant_op.constant([0.02], shape=[1, 1], dtype=dtype),
+              constant_op.constant([1]), constant_op.constant([2, 1]))
+
+          opt = ftrl.FtrlOptimizer(
+              3.0,
+              initial_accumulator_value=0.1,
+              l1_regularization_strength=0.001,
+              l2_regularization_strength=2.0,
+              l2_shrinkage_regularization_strength=0.1)
+          update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+          self.evaluate(variables.global_variables_initializer())
+
+          v0_val, v1_val = self.evaluate([var0, var1])
+          self.assertAllCloseAccordingToType([[1.0], [2.0]], v0_val)
+          self.assertAllCloseAccordingToType([[4.0], [3.0]], v1_val)
+
+          # Run 10 steps FTRL
+          for _ in range(10):
+            update.run()
+
+          v0_val, v1_val = self.evaluate([var0, var1])
+          self.assertAllCloseAccordingToType([[-0.22578995], [2.]], v0_val)
+          self.assertAllCloseAccordingToType([[4.], [-0.13229476]], v1_val)
+
+  def testFtrlWithL2ShrinkageDoesNotChangeLrSchedule(self):
+    """Verifies that l2 shrinkage in FTRL does not change lr schedule."""
+    # The v1 optimizers do not support eager execution
+    with ops.Graph().as_default():
+      for dtype in [dtypes.half, dtypes.float32]:
+        with self.cached_session():
+          var0 = variables.Variable([1.0, 2.0], dtype=dtype)
+          var1 = variables.Variable([1.0, 2.0], dtype=dtype)
+          grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+          grads1 = constant_op.constant([0.1, 0.2], dtype=dtype)
+
+          opt0 = ftrl.FtrlOptimizer(
+              3.0,
+              initial_accumulator_value=0.1,
+              l1_regularization_strength=0.001,
+              l2_regularization_strength=2.0,
+              l2_shrinkage_regularization_strength=0.1)
+          opt1 = ftrl.FtrlOptimizer(
+              3.0,
+              initial_accumulator_value=0.1,
+              l1_regularization_strength=0.001,
+              l2_regularization_strength=2.0)
+          update0 = opt0.apply_gradients([(grads0, var0)])
+          update1 = opt1.apply_gradients([(grads1, var1)])
+          self.evaluate(variables.global_variables_initializer())
+
+          v0_val, v1_val = self.evaluate([var0, var1])
+          self.assertAllCloseAccordingToType([1.0, 2.0], v0_val)
+          self.assertAllCloseAccordingToType([1.0, 2.0], v1_val)
+
+          # Run 10 steps FTRL
+          for _ in range(10):
+            update0.run()
+            update1.run()
+
+          v0_val, v1_val = self.evaluate([var0, var1])
+          # var0 is experiencing L2 shrinkage so it should be smaller than var1
+          # in magnitude.
+          self.assertTrue((v0_val**2 < v1_val**2).all())
+          accum0 = list(self.evaluate(opt0._slots)["accum"].values())[0]
+          accum1 = list(self.evaluate(opt1._slots)["accum"].values())[0]
+          # L2 shrinkage should not change how we update grad accumulator.
+          self.assertAllCloseAccordingToType(accum0, accum1)
+
+  def applyOptimizer(self, opt, dtype, steps=5, is_sparse=False):
+    if is_sparse:
+      var0 = variables.Variable([[0.0], [0.0]], dtype=dtype)
+      var1 = variables.Variable([[0.0], [0.0]], dtype=dtype)
+      grads0 = indexed_slices.IndexedSlices(
+          constant_op.constant([0.1], shape=[1, 1], dtype=dtype),
+          constant_op.constant([0]), constant_op.constant([2, 1]))
+      grads1 = indexed_slices.IndexedSlices(
+          constant_op.constant([0.02], shape=[1, 1], dtype=dtype),
+          constant_op.constant([1]), constant_op.constant([2, 1]))
+    else:
+      var0 = variables.Variable([0.0, 0.0], dtype=dtype)
+      var1 = variables.Variable([0.0, 0.0], dtype=dtype)
+      grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+      grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+
+    update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+    self.evaluate(variables.global_variables_initializer())
+
+    sess = ops.get_default_session()
+    v0_val, v1_val = self.evaluate([var0, var1])
+    if is_sparse:
+      self.assertAllCloseAccordingToType([[0.0], [0.0]], v0_val)
+      self.assertAllCloseAccordingToType([[0.0], [0.0]], v1_val)
+    else:
+      self.assertAllCloseAccordingToType([0.0, 0.0], v0_val)
+      self.assertAllCloseAccordingToType([0.0, 0.0], v1_val)
+
+    # Run Ftrl for a few steps
+    for _ in range(steps):
+      update.run()
+
+    v0_val, v1_val = self.evaluate([var0, var1])
+    return v0_val, v1_val
+
+  # When variables are initialized with Zero, FTRL-Proximal has two properties:
+  # 1. Without L1&L2 but with fixed learning rate, FTRL-Proximal is identical
+  # with GradientDescent.
+  # 2. Without L1&L2 but with adaptive learning rate, FTRL-Proximal is identical
+  # with Adagrad.
+  # So, basing on these two properties, we test if our implementation of
+  # FTRL-Proximal performs same updates as Adagrad or GradientDescent.
+  def testEquivAdagradwithoutRegularization(self):
+    # The v1 optimizers do not support eager execution
+    with ops.Graph().as_default():
+      for dtype in [dtypes.half, dtypes.float32]:
+        with self.cached_session():
+          val0, val1 = self.applyOptimizer(
+              ftrl.FtrlOptimizer(
+                  3.0,
+                  # Adagrad learning rate
+                  learning_rate_power=-0.5,
+                  initial_accumulator_value=0.1,
+                  l1_regularization_strength=0.0,
+                  l2_regularization_strength=0.0),
+              dtype)
+
+        with self.cached_session():
+          val2, val3 = self.applyOptimizer(
+              adagrad.AdagradOptimizer(3.0, initial_accumulator_value=0.1),
+              dtype)
+
+        self.assertAllCloseAccordingToType(val0, val2, half_rtol=2e-3)
+        self.assertAllCloseAccordingToType(val1, val3, half_rtol=2e-3)
+
+  def testEquivSparseAdagradwithoutRegularization(self):
+    # The v1 optimizers do not support eager execution
+    with ops.Graph().as_default():
+      for dtype in [dtypes.half, dtypes.float32]:
+        with self.cached_session():
+          val0, val1 = self.applyOptimizer(
+              ftrl.FtrlOptimizer(
+                  3.0,
+                  # Adagrad learning rate
+                  learning_rate_power=-0.5,
+                  initial_accumulator_value=0.1,
+                  l1_regularization_strength=0.0,
+                  l2_regularization_strength=0.0),
+              dtype,
+              is_sparse=True)
+
+        with self.cached_session():
+          val2, val3 = self.applyOptimizer(
+              adagrad.AdagradOptimizer(3.0, initial_accumulator_value=0.1),
+              dtype,
+              is_sparse=True)
+
+        self.assertAllCloseAccordingToType(val0, val2)
+        self.assertAllCloseAccordingToType(val1, val3)
+
+  def testEquivSparseGradientDescentwithoutRegularization(self):
+    # The v1 optimizers do not support eager execution
+    with ops.Graph().as_default():
+      for dtype in [dtypes.half, dtypes.float32]:
+        with self.cached_session():
+          val0, val1 = self.applyOptimizer(
+              ftrl.FtrlOptimizer(
+                  3.0,
+                  # Fixed learning rate
+                  learning_rate_power=-0.0,
+                  initial_accumulator_value=0.1,
+                  l1_regularization_strength=0.0,
+                  l2_regularization_strength=0.0),
+              dtype,
+              is_sparse=True)
+
+        with self.cached_session():
+          val2, val3 = self.applyOptimizer(
+              gradient_descent.GradientDescentOptimizer(3.0),
+              dtype,
+              is_sparse=True)
+
+        self.assertAllCloseAccordingToType(val0, val2)
+        self.assertAllCloseAccordingToType(val1, val3)
+
+  def testEquivGradientDescentwithoutRegularization(self):
+    # The v1 optimizers do not support eager execution
+    with ops.Graph().as_default():
+      for dtype in [dtypes.half, dtypes.float32]:
+        with self.cached_session():
+          val0, val1 = self.applyOptimizer(
+              ftrl.FtrlOptimizer(
+                  3.0,
+                  # Fixed learning rate
+                  learning_rate_power=-0.0,
+                  initial_accumulator_value=0.1,
+                  l1_regularization_strength=0.0,
+                  l2_regularization_strength=0.0),
+              dtype)
+
+        with self.cached_session():
+          val2, val3 = self.applyOptimizer(
+              gradient_descent.GradientDescentOptimizer(3.0), dtype)
+
+        self.assertAllCloseAccordingToType(val0, val2)
+        self.assertAllCloseAccordingToType(val1, val3)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/gradient_descent.py

@@ -0,0 +1,82 @@
+# 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.
+# ==============================================================================
+
+"""GradientDescent for TensorFlow."""
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import gen_training_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.training import optimizer
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(v1=["train.GradientDescentOptimizer"])
+class GradientDescentOptimizer(optimizer.Optimizer):
+  """Optimizer that implements the gradient descent algorithm.
+  """
+
+  def __init__(self, learning_rate, use_locking=False, name="GradientDescent"):
+    """Construct a new gradient descent optimizer.
+
+    Args:
+      learning_rate: A Tensor or a floating point value.  The learning
+        rate to use.
+      use_locking: If True use locks for update operations.
+      name: Optional name prefix for the operations created when applying
+        gradients. Defaults to "GradientDescent".
+
+    @compatibility(eager)
+    When eager execution is enabled, `learning_rate` can 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(GradientDescentOptimizer, self).__init__(use_locking, name)
+    self._learning_rate = learning_rate
+    self._learning_rate_tensor = None
+
+  def _apply_dense(self, grad, var):
+    return gen_training_ops.apply_gradient_descent(
+        var,
+        math_ops.cast(self._learning_rate_tensor, var.dtype.base_dtype),
+        grad,
+        use_locking=self._use_locking).op
+
+  def _resource_apply_dense(self, grad, handle):
+    return gen_training_ops.resource_apply_gradient_descent(
+        handle.handle, math_ops.cast(self._learning_rate_tensor,
+                                     grad.dtype.base_dtype),
+        grad, use_locking=self._use_locking)
+
+  def _resource_apply_sparse_duplicate_indices(self, grad, handle, indices):
+    return resource_variable_ops.resource_scatter_add(
+        handle.handle,
+        indices,
+        -grad * math_ops.cast(self._learning_rate_tensor,
+                              grad.dtype.base_dtype))
+
+  def _apply_sparse_duplicate_indices(self, grad, var):
+    delta = indexed_slices.IndexedSlices(
+        grad.values *
+        math_ops.cast(self._learning_rate_tensor, var.dtype.base_dtype),
+        grad.indices, grad.dense_shape)
+    return var.scatter_sub(delta, use_locking=self._use_locking)
+
+  def _prepare(self):
+    learning_rate = self._call_if_callable(self._learning_rate)
+    self._learning_rate_tensor = ops.convert_to_tensor(
+        learning_rate, name="learning_rate")

tensorflow/python/training/gradient_descent_test.py

@@ -0,0 +1,255 @@
+# 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 tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import embedding_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import resources
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import gradient_descent
+
+
+class GradientDescentOptimizerTest(test.TestCase):
+
+  def testBasic(self):
+    for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+      # train.GradientDescentOptimizer is V1 only API.
+      with ops.Graph().as_default(), self.cached_session():
+        var0 = variables.Variable([1.0, 2.0], dtype=dtype)
+        var1 = variables.Variable([3.0, 4.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.1], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.01], dtype=dtype)
+        optimizer = gradient_descent.GradientDescentOptimizer(3.0)
+        sgd_op = optimizer.apply_gradients(
+            zip([grads0, grads1], [var0, var1]))
+        self.evaluate(variables.global_variables_initializer())
+        # Fetch params to validate initial values
+        self.assertAllCloseAccordingToType([1.0, 2.0], self.evaluate(var0))
+        self.assertAllCloseAccordingToType([3.0, 4.0], self.evaluate(var1))
+        # Run 1 step of sgd
+        sgd_op.run()
+        # 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))
+        self.assertEqual(0, len(optimizer.variables()))
+
+  def testBasicResourceVariable(self):
+    for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+      # train.GradientDescentOptimizer is V1 only API.
+      with ops.Graph().as_default(), self.cached_session():
+        var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
+        var1 = resource_variable_ops.ResourceVariable([3.0, 4.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.1], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.01], dtype=dtype)
+        sgd_op = gradient_descent.GradientDescentOptimizer(3.0).apply_gradients(
+            zip([grads0, grads1], [var0, var1]))
+        # TODO(apassos) calling initialize_resources on all resources here
+        # doesn't work because the sessions and graph are reused across unit
+        # tests and this would mean trying to reinitialize variables. Figure out
+        # a long-term solution for this.
+        resources.initialize_resources([var0, var1]).run()
+        # Fetch params to validate initial values
+        self.assertAllCloseAccordingToType([1.0, 2.0], self.evaluate(var0))
+        self.assertAllCloseAccordingToType([3.0, 4.0], self.evaluate(var1))
+        # Run 1 step of sgd
+        sgd_op.run()
+        # 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 testBasicCallableParams(self):
+    for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+      # train.GradientDescentOptimizer is V1 only API.
+      with ops.Graph().as_default(), self.cached_session():
+        var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
+        var1 = resource_variable_ops.ResourceVariable([3.0, 4.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.1], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.01], dtype=dtype)
+        lr = lambda: 3.0
+        sgd_op = gradient_descent.GradientDescentOptimizer(lr).apply_gradients(
+            zip([grads0, grads1], [var0, var1]))
+        # TODO(apassos) calling initialize_resources on all resources here
+        # doesn't work because the sessions and graph are reused across unit
+        # tests and this would mean trying to reinitialize variables. Figure out
+        # a long-term solution for this.
+        resources.initialize_resources([var0, var1]).run()
+        # Fetch params to validate initial values
+        self.assertAllCloseAccordingToType([1.0, 2.0], self.evaluate(var0))
+        self.assertAllCloseAccordingToType([3.0, 4.0], self.evaluate(var1))
+        # Run 1 step of sgd
+        sgd_op.run()
+        # 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 testMinimizeResourceVariable(self):
+    for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+      # train.GradientDescentOptimizer is V1 only API.
+      with ops.Graph().as_default(), self.cached_session():
+        var0 = resource_variable_ops.ResourceVariable([[1.0, 2.0]], dtype=dtype)
+        var1 = resource_variable_ops.ResourceVariable([3.0], dtype=dtype)
+        x = constant_op.constant([[4.0], [5.0]], dtype=dtype)
+        pred = math_ops.matmul(var0, x) + var1
+        loss = pred * pred
+        sgd_op = gradient_descent.GradientDescentOptimizer(1.0).minimize(loss)
+        # TODO(apassos) calling initialize_resources on all resources here
+        # doesn't work because the sessions and graph are reused across unit
+        # tests and this would mean trying to reinitialize variables. Figure out
+        # a long-term solution for this.
+        resources.initialize_resources([var0, var1]).run()
+        # Fetch params to validate initial values
+        self.assertAllCloseAccordingToType([[1.0, 2.0]], self.evaluate(var0))
+        self.assertAllCloseAccordingToType([3.0], self.evaluate(var1))
+        # Run 1 step of sgd
+        sgd_op.run()
+        # 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 testMinimizeSparseResourceVariable(self):
+    for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+      # train.GradientDescentOptimizer is V1 only API.
+      with ops.Graph().as_default(), self.cached_session():
+        var0 = resource_variable_ops.ResourceVariable([[1.0, 2.0]], dtype=dtype)
+        var1 = resource_variable_ops.ResourceVariable([3.0], dtype=dtype)
+        x = constant_op.constant([[4.0], [5.0]], dtype=dtype)
+        pred = math_ops.matmul(embedding_ops.embedding_lookup([var0], [0]), x)
+        pred += var1
+        loss = pred * pred
+        sgd_op = gradient_descent.GradientDescentOptimizer(1.0).minimize(loss)
+        # TODO(apassos) calling initialize_resources on all resources here
+        # doesn't work because the sessions and graph are reused across unit
+        # tests and this would mean trying to reinitialize variables. Figure out
+        # a long-term solution for this.
+        self.evaluate(variables.global_variables_initializer())
+        # Fetch params to validate initial values
+        self.assertAllCloseAccordingToType([[1.0, 2.0]], self.evaluate(var0))
+        self.assertAllCloseAccordingToType([3.0], self.evaluate(var1))
+        # Run 1 step of sgd
+        sgd_op.run()
+        # 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 [dtypes.half, dtypes.float32, dtypes.float64]:
+      # train.GradientDescentOptimizer is V1 only API.
+      with ops.Graph().as_default(), self.cached_session():
+        var0 = variables.Variable([1.0, 2.0], dtype=dtype)
+        var1 = variables.Variable([3.0, 4.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.1], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.01], dtype=dtype)
+        lrate = constant_op.constant(3.0)
+        sgd_op = gradient_descent.GradientDescentOptimizer(
+            lrate).apply_gradients(zip([grads0, grads1], [var0, var1]))
+        self.evaluate(variables.global_variables_initializer())
+        # Fetch params to validate initial values
+        self.assertAllCloseAccordingToType([1.0, 2.0], self.evaluate(var0))
+        self.assertAllCloseAccordingToType([3.0, 4.0], self.evaluate(var1))
+        # Run 1 step of sgd
+        sgd_op.run()
+        # 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):
+    for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+      # train.GradientDescentOptimizer is V1 only API.
+      with ops.Graph().as_default(), self.cached_session():
+        opt = gradient_descent.GradientDescentOptimizer(3.0)
+        values = [1.0, 3.0]
+        vars_ = [variables.Variable([v], dtype=dtype) for v in values]
+        grads_and_vars = opt.compute_gradients(vars_[0] + vars_[1], vars_)
+        self.evaluate(variables.global_variables_initializer())
+        for grad, _ in grads_and_vars:
+          self.assertAllCloseAccordingToType([1.0], self.evaluate(grad))
+
+  def testWithGlobalStep(self):
+    for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+      # train.GradientDescentOptimizer is V1 only API.
+      with ops.Graph().as_default(), self.cached_session():
+        global_step = variables.Variable(0, trainable=False)
+        var0 = variables.Variable([1.0, 2.0], dtype=dtype)
+        var1 = variables.Variable([3.0, 4.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.1], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.01], dtype=dtype)
+        sgd_op = gradient_descent.GradientDescentOptimizer(3.0).apply_gradients(
+            zip([grads0, grads1], [var0, var1]), global_step=global_step)
+        self.evaluate(variables.global_variables_initializer())
+        # Fetch params to validate initial values
+        self.assertAllCloseAccordingToType([1.0, 2.0], self.evaluate(var0))
+        self.assertAllCloseAccordingToType([3.0, 4.0], self.evaluate(var1))
+        # Run 1 step of sgd
+        sgd_op.run()
+        # Validate updated params and global_step
+        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))
+        self.assertAllCloseAccordingToType(1, self.evaluate(global_step))
+
+  def testSparseBasic(self):
+    for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+      # train.GradientDescentOptimizer is V1 only API.
+      with ops.Graph().as_default(), self.cached_session():
+        var0 = variables.Variable([[1.0], [2.0]], dtype=dtype)
+        var1 = variables.Variable([[3.0], [4.0]], dtype=dtype)
+        grads0 = indexed_slices.IndexedSlices(
+            constant_op.constant(
+                [0.1], shape=[1, 1], dtype=dtype),
+            constant_op.constant([0]),
+            constant_op.constant([2, 1]))
+        grads1 = indexed_slices.IndexedSlices(
+            constant_op.constant(
+                [0.01], shape=[1, 1], dtype=dtype),
+            constant_op.constant([1]),
+            constant_op.constant([2, 1]))
+        sgd_op = gradient_descent.GradientDescentOptimizer(3.0).apply_gradients(
+            zip([grads0, grads1], [var0, var1]))
+        self.evaluate(variables.global_variables_initializer())
+        # Fetch params to validate initial values
+        self.assertAllCloseAccordingToType([[1.0], [2.0]], self.evaluate(var0))
+        self.assertAllCloseAccordingToType([[3.0], [4.0]], self.evaluate(var1))
+        # Run 1 step of sgd
+        sgd_op.run()
+        # 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))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/learning_rate_decay.py

@@ -0,0 +1,28 @@
+# 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."""
+
+from tensorflow.python.keras.optimizer_v2 import legacy_learning_rate_decay as learning_rate_decay
+
+
+exponential_decay = learning_rate_decay.exponential_decay
+piecewise_constant = learning_rate_decay.piecewise_constant
+polynomial_decay = learning_rate_decay.polynomial_decay
+natural_exp_decay = learning_rate_decay.natural_exp_decay
+inverse_time_decay = learning_rate_decay.inverse_time_decay
+cosine_decay = learning_rate_decay.cosine_decay
+cosine_decay_restarts = learning_rate_decay.cosine_decay_restarts
+linear_cosine_decay = learning_rate_decay.linear_cosine_decay
+noisy_linear_cosine_decay = learning_rate_decay.noisy_linear_cosine_decay

tensorflow/python/training/localhost_cluster_performance_test.py

@@ -0,0 +1,115 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Tests and benchmarks for creating RPC clusters on localhost."""
+
+import time
+
+import numpy as np
+
+from tensorflow.python.client import session as session_lib
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import partitioned_variables
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import device_setter
+
+
+class CreateLocalClusterTest(test.TestCase):
+
+  @test_util.run_v1_only("b/120545219")
+  def testCreateLocalCluster(self):
+    workers, _ = test.create_local_cluster(num_workers=2, num_ps=2)
+    worker_sessions = [session_lib.Session(w.target) for w in workers]
+    with ops.device("/job:ps/task:0"):
+      var0 = variables.Variable(0.0)
+    with ops.device("/job:ps/task:1"):
+      var1 = variables.Variable(1.0)
+    worker_sessions[0].run([var0.initializer, var1.initializer])
+    with ops.device("/job:ps/task:0"):
+      var2 = variables.Variable(2.0)
+    with ops.device("/job:ps/task:1"):
+      var3 = variables.Variable(3.0)
+    worker_sessions[1].run([var2.initializer, var3.initializer])
+
+    # Read values back in the opposite session
+    self.assertAllEqual(0.0, var0.eval(session=worker_sessions[1]))
+    self.assertAllEqual(1.0, var1.eval(session=worker_sessions[1]))
+    self.assertAllEqual(2.0, var2.eval(session=worker_sessions[0]))
+    self.assertAllEqual(3.0, var3.eval(session=worker_sessions[0]))
+
+
+class CreateLocalClusterBenchmark(test.Benchmark):
+
+  def benchmarkCreateLocalCluster(self):
+    deltas = []
+    iters = 5
+    for _ in range(iters):
+      start_time = time.time()
+      test.create_local_cluster(num_workers=1, num_ps=10)
+      end_time = time.time()
+      deltas.append(end_time - start_time)
+
+    median_deltas = np.median(deltas)
+    print("\n\nbenchmark_create_local_cluster_1_worker_10_ps.  "
+          "iterations: %d, median wall time: %g\n\n" % (iters, median_deltas))
+    self.report_benchmark(
+        iters=iters,
+        wall_time=median_deltas,
+        name="benchmark_create_local_cluster_1_worker_10_ps")
+
+
+class PartitionedVariablesBenchmark(test.Benchmark):
+
+  def benchmark_create_1000_partitions_with_100_parameter_servers(self):
+    workers, _ = test.create_local_cluster(num_workers=1, num_ps=100)
+    worker_sessions = [session_lib.Session(w.target) for w in workers]
+    worker = worker_sessions[0]
+    partition_sizes = (1, 512, 1024 * 32, 1024 * 128)
+
+    partitioned = []
+
+    for partition_size in partition_sizes:
+      # max_shard_bytes is 4, shape is 1000*partition_size float32s which should
+      # partition into 1000 shards, each containing partition_size float32s.
+      print("Building partitioned variable with %d floats per partition" %
+            partition_size)
+      with ops.device(device_setter.replica_device_setter(ps_tasks=100)):
+        partitioned_ix = variable_scope.get_variable(
+            "partitioned_%d" % partition_size,
+            shape=[1000 * partition_size],
+            dtype=dtypes.float32,
+            # Each partition to have exactly N float32s
+            partitioner=partitioned_variables.variable_axis_size_partitioner(
+                max_shard_bytes=4 * partition_size))
+        # Concatenates along axis 0
+        partitioned.append(ops.convert_to_tensor(partitioned_ix))
+
+    variables.global_variables_initializer().run(session=worker)
+
+    for ix, partition_size in enumerate(partition_sizes):
+      print("Running benchmark having partitions with %d floats" %
+            partition_size)
+      self.run_op_benchmark(
+          worker,
+          partitioned[ix],
+          name=("read_concat_1000_partitions_from_"
+                "100_parameter_servers_partsize_%d_floats" % partition_size))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/momentum.py

@@ -0,0 +1,203 @@
+# 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.
+# ==============================================================================
+
+"""Momentum for TensorFlow."""
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import gen_training_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.training import optimizer
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(v1=["train.MomentumOptimizer"])
+class MomentumOptimizer(optimizer.Optimizer):
+  """Optimizer that implements the Momentum algorithm.
+
+  Computes (if `use_nesterov = False`):
+
+  ```
+  accumulation = momentum * accumulation + gradient
+  variable -= learning_rate * accumulation
+  ```
+
+  Note that in the dense version of this algorithm, `accumulation` is updated
+  and applied regardless of a gradient's value, whereas the sparse version (when
+  the gradient is an `IndexedSlices`, typically because of `tf.gather` or an
+  embedding) only updates variable slices and corresponding `accumulation` terms
+  when that part of the variable was used in the forward pass.
+
+  @compatibility(TF2)
+  tf.compat.v1.train.MomentumOptimizer is compatible with eager mode and
+  `tf.function`.
+  When eager execution is enabled, `learning_rate`,`momentum`, 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.
+
+  To switch to native TF2 style, please directly use
+  [`tf.keras.optimizers.SGD`]
+  (https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/SGD)
+  with the `momentum` argument.
+
+  #### Structural mapping to native TF2
+
+  Before:
+
+  ```python
+  optimizer = tf.compat.v1.train.MomentumOptimizer(
+    learning_rate=learning_rate,
+    momentum=momentum,
+    use_nesterov=use_nesterov)
+  ```
+
+  After:
+
+  ```python
+  optimizer = tf.keras.optimizers.SGD(
+    learning_rate=learning_rate,
+    momentum=momentum,
+    nesterov=use_nesterov)
+  ```
+
+  #### How to map arguments
+  | TF1 Arg Name       | TF2 Arg Name   | Note                             |
+  | ------------------ | -------------  | -------------------------------  |
+  | `learning_rate`    | `learning_rate`| Be careful of setting           |
+  : : : learning_rate tensor value computed from the global step.          :
+  : : : In TF1 this was usually meant to imply a dynamic learning rate and :
+  : : : would recompute in each step. In TF2 (eager + function) it will    :
+  : : : treat it as a scalar value that only gets computed once instead of :
+  : : : a symbolic placeholder to be computed each time.                   :
+  | `momentum`         | `momentum`     | -                                |
+  | `use_locking`      | -              | Not applicable in TF2.           |
+  | `use_nesterov`     | `nesterov`     | -                                |
+
+  #### Before & after usage example
+  Before:
+
+  ```python
+  x = tf.Variable([1,2,3], dtype=tf.float32)
+  grad = tf.constant([0.1, 0.2, 0.3])
+  optimizer = tf.compat.v1.train.MomentumOptimizer(
+    learning_rate=0.001,
+    momentum=0.9,
+    use_nesterov=False)
+  optimizer.apply_gradients(zip([grad], [x]))
+  ```
+
+  After:
+
+  ```python
+  x = tf.Variable([1,2,3], dtype=tf.float32)
+  grad = tf.constant([0.1, 0.2, 0.3])
+  optimizer = tf.keras.optimizers.SGD(
+    learning_rate=0.001,
+    momentum=0.9,
+    nesterov=False)
+  optimizer.apply_gradients(zip([grad], [x]))
+  ```
+
+  @end_compatibility
+
+  """
+
+  def __init__(self, learning_rate, momentum,
+               use_locking=False, name="Momentum", use_nesterov=False):
+    """Construct a new Momentum optimizer.
+
+    Args:
+      learning_rate: A `Tensor` or a floating point value.  The learning rate.
+      momentum: A `Tensor` or a floating point value.  The momentum.
+      use_locking: If `True` use locks for update operations.
+      name: Optional name prefix for the operations created when applying
+        gradients.  Defaults to "Momentum".
+      use_nesterov: If `True` use Nesterov Momentum.
+        See (Sutskever et al., 2013).
+        This implementation always computes gradients at the value of the
+        variable(s) passed to the optimizer. Using Nesterov Momentum makes the
+        variable(s) track the values called `theta_t + mu*v_t` in the paper.
+        This implementation is an approximation of the original formula, valid
+        for high values of momentum. It will compute the "adjusted gradient"
+        in NAG by assuming that the new gradient will be estimated by the
+        current average gradient plus the product of momentum and the change
+        in the average gradient.
+
+    References:
+      On the importance of initialization and momentum in deep learning:
+        [Sutskever et al., 2013]
+        (http://proceedings.mlr.press/v28/sutskever13.html)
+        ([pdf](http://proceedings.mlr.press/v28/sutskever13.pdf))
+
+
+    """
+    super(MomentumOptimizer, self).__init__(use_locking, name)
+    self._learning_rate = learning_rate
+    self._momentum = momentum
+    self._use_nesterov = use_nesterov
+
+  def _create_slots(self, var_list):
+    for v in var_list:
+      self._zeros_slot(v, "momentum", self._name)
+
+  def _prepare(self):
+    learning_rate = self._learning_rate
+    if callable(learning_rate):
+      learning_rate = learning_rate()
+    self._learning_rate_tensor = ops.convert_to_tensor(learning_rate,
+                                                       name="learning_rate")
+    momentum = self._momentum
+    if callable(momentum):
+      momentum = momentum()
+    self._momentum_tensor = ops.convert_to_tensor(momentum, name="momentum")
+
+  def _apply_dense(self, grad, var):
+    mom = self.get_slot(var, "momentum")
+    return gen_training_ops.apply_momentum(
+        var, mom,
+        math_ops.cast(self._learning_rate_tensor, var.dtype.base_dtype),
+        grad,
+        math_ops.cast(self._momentum_tensor, var.dtype.base_dtype),
+        use_locking=self._use_locking,
+        use_nesterov=self._use_nesterov).op
+
+  def _resource_apply_dense(self, grad, var):
+    mom = self.get_slot(var, "momentum")
+    return gen_training_ops.resource_apply_momentum(
+        var.handle, mom.handle,
+        math_ops.cast(self._learning_rate_tensor, grad.dtype.base_dtype),
+        grad,
+        math_ops.cast(self._momentum_tensor, grad.dtype.base_dtype),
+        use_locking=self._use_locking,
+        use_nesterov=self._use_nesterov)
+
+  def _apply_sparse(self, grad, var):
+    mom = self.get_slot(var, "momentum")
+    return gen_training_ops.sparse_apply_momentum(
+        var, mom,
+        math_ops.cast(self._learning_rate_tensor, var.dtype.base_dtype),
+        grad.values, grad.indices,
+        math_ops.cast(self._momentum_tensor, var.dtype.base_dtype),
+        use_locking=self._use_locking,
+        use_nesterov=self._use_nesterov).op
+
+  def _resource_apply_sparse(self, grad, var, indices):
+    mom = self.get_slot(var, "momentum")
+    return gen_training_ops.resource_sparse_apply_momentum(
+        var.handle, mom.handle,
+        math_ops.cast(self._learning_rate_tensor, grad.dtype),
+        grad, indices,
+        math_ops.cast(self._momentum_tensor, grad.dtype),
+        use_locking=self._use_locking,
+        use_nesterov=self._use_nesterov)

tensorflow/python/training/momentum_test.py

@@ -0,0 +1,599 @@
+# 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.
+# ==============================================================================
+"""Tests for Momentum."""
+
+import numpy as np
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import embedding_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import momentum as momentum_lib
+
+
+class MomentumOptimizerTest(test.TestCase):
+
+  def _update_nesterov_momentum_numpy(self, var, accum, g, lr, momentum):
+    var = var + accum * lr * momentum
+    accum = accum * momentum + g
+    var = var - lr * accum
+    var = var - accum * lr * momentum
+    return var, accum
+
+  def doTestBasic(self, use_resource=False, use_callable_params=False):
+    for i, dtype in enumerate([dtypes.half, dtypes.float32, dtypes.float64]):
+      if use_resource:
+        var0 = resource_variable_ops.ResourceVariable([1.0, 2.0],
+                                                      dtype=dtype,
+                                                      name="var0_%d" % i)
+        var1 = resource_variable_ops.ResourceVariable([3.0, 4.0],
+                                                      dtype=dtype,
+                                                      name="var1_%d" % i)
+      else:
+        var0 = variables.Variable([1.0, 2.0], dtype=dtype)
+        var1 = variables.Variable([3.0, 4.0], dtype=dtype)
+      grads0 = constant_op.constant([0.1, 0.1], dtype=dtype)
+      grads1 = constant_op.constant([0.01, 0.01], dtype=dtype)
+      learning_rate = lambda: 2.0
+      momentum = lambda: 0.9
+      if not use_callable_params:
+        learning_rate = learning_rate()
+        momentum = momentum()
+      mom_opt = momentum_lib.MomentumOptimizer(
+          learning_rate=learning_rate, momentum=momentum)
+      mom_update = mom_opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+
+      if not context.executing_eagerly():
+        self.evaluate(variables.global_variables_initializer())
+        # Fetch params to validate initial values
+        self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+        self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+
+      # Check we have slots
+      self.assertEqual(["momentum"], mom_opt.get_slot_names())
+      slot0 = mom_opt.get_slot(var0, "momentum")
+      self.assertEqual(slot0.get_shape(), var0.get_shape())
+      slot1 = mom_opt.get_slot(var1, "momentum")
+      self.assertEqual(slot1.get_shape(), var1.get_shape())
+      if not context.executing_eagerly():
+        self.assertFalse(slot0 in variables.trainable_variables())
+        self.assertFalse(slot1 in variables.trainable_variables())
+
+      # Step 1: the momentum accumulators where 0. So we should see a normal
+      # update: v -= grad * learning_rate
+      if not context.executing_eagerly():
+        self.evaluate(mom_update)
+      # Check that the momentum accumulators have been updated.
+      self.assertAllCloseAccordingToType(
+          np.array([0.1, 0.1]), self.evaluate(slot0))
+      self.assertAllCloseAccordingToType(
+          np.array([0.01, 0.01]), 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.
+      if context.executing_eagerly():
+        mom_opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+      else:
+        self.evaluate(mom_update)
+      # Check that the momentum accumulators have been updated.
+      self.assertAllCloseAccordingToType(
+          np.array([(0.9 * 0.1 + 0.1), (0.9 * 0.1 + 0.1)]),
+          self.evaluate(slot0))
+      self.assertAllCloseAccordingToType(
+          np.array([(0.9 * 0.01 + 0.01), (0.9 * 0.01 + 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 testBasic(self):
+    with self.cached_session():
+      self.doTestBasic(use_resource=False)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testResourceBasic(self):
+    self.doTestBasic(use_resource=True)
+
+  def testBasicCallableParams(self):
+    with context.eager_mode():
+      self.doTestBasic(use_resource=True, use_callable_params=True)
+
+  def testVariablesAcrossGraphs(self):
+    optimizer = momentum_lib.MomentumOptimizer(0.01, 0.5)
+    with ops.Graph().as_default():
+      var0 = resource_variable_ops.ResourceVariable([1.0, 2.0],
+                                                    dtype=dtypes.float32,
+                                                    name="var0")
+      var1 = resource_variable_ops.ResourceVariable([3.0, 4.0],
+                                                    dtype=dtypes.float32,
+                                                    name="var1")
+      loss = math_ops.reduce_sum(var0 + var1)
+      optimizer.minimize(loss)
+      optimizer_variables = optimizer.variables()
+      self.assertStartsWith(optimizer_variables[0].name, "var0")
+      self.assertStartsWith(optimizer_variables[1].name, "var1")
+      self.assertEqual(2, len(optimizer_variables))
+
+    with ops.Graph().as_default():
+      var2 = resource_variable_ops.ResourceVariable([1.0, 2.0],
+                                                    dtype=dtypes.float32,
+                                                    name="var2")
+      var3 = resource_variable_ops.ResourceVariable([3.0, 4.0],
+                                                    dtype=dtypes.float32,
+                                                    name="var3")
+      loss = math_ops.reduce_sum(var2 + var3)
+      optimizer.minimize(loss)
+      optimizer_variables = optimizer.variables()
+      self.assertStartsWith(optimizer_variables[0].name, "var2")
+      self.assertStartsWith(optimizer_variables[1].name, "var3")
+      self.assertEqual(2, len(optimizer_variables))
+
+  def testNesterovMomentum(self):
+    for dtype in [dtypes.float32, dtypes.float64]:
+      # train.MomentumOptimizer is V1 only API.
+      with ops.Graph().as_default(), self.cached_session():
+        var0 = variables.Variable([1.0, 2.0], dtype=dtype)
+        var1 = variables.Variable([3.0, 4.0], dtype=dtype)
+        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)
+        cost = 5 * var0 * var0 + 3 * var1
+        global_step = variables.Variable(
+            array_ops.zeros([], dtypes.int64), name="global_step")
+        mom_op = momentum_lib.MomentumOptimizer(
+            learning_rate=2.0, momentum=0.9, use_nesterov=True)
+        opt_op = mom_op.minimize(cost, global_step, [var0, var1])
+        self.evaluate(variables.global_variables_initializer())
+        for t in range(1, 5):
+          opt_op.run()
+          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):
+    for dtype in [dtypes.float32, dtypes.float64]:
+      # train.MomentumOptimizer is V1 only API.
+      with ops.Graph().as_default(), self.cached_session():
+        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 = variables.Variable(var0_np)
+        var1 = variables.Variable(var1_np)
+        mom_op = momentum_lib.MomentumOptimizer(
+            learning_rate=2.0, momentum=0.9, use_nesterov=True)
+        x_feed = array_ops.placeholder(dtype)
+        y_feed = indexed_slices.IndexedSlices(x_feed,
+                                              constant_op.constant([0, 1]),
+                                              constant_op.constant([2]))
+        grads_and_vars = [(y_feed, var0),
+                          (constant_op.constant([3.0, 3.0], dtype=dtype), var1)]
+        opt_update = mom_op.apply_gradients(grads_and_vars)
+        self.evaluate(variables.global_variables_initializer())
+        for t in range(1, 5):
+          opt_update.run(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))
+
+  @test_util.run_in_graph_and_eager_modes
+  def testMinimizeSparseResourceVariable(self):
+    for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+      # This test invokes the ResourceSparseApplyMomentum operation, which
+      # did not have a registered GPU kernel as of April 2018. With graph
+      # execution, the placement algorithm notices this and automatically
+      # places the variable in CPU (host) memory. With eager execution,
+      # the variable would be placed in GPU memory if available, which
+      # would then conflict with the future invocation of the
+      # ResourceSparseApplyMomentum operation.
+      # To work around this discrepancy, for now we force the variable
+      # to be placed on CPU.
+      with ops.device("/cpu:0"):
+        var0 = resource_variable_ops.ResourceVariable([[1.0, 2.0]], dtype=dtype)
+
+      # pylint: disable=cell-var-from-loop
+      def loss():
+        x = constant_op.constant([[4.0], [5.0]], dtype=dtype)
+        pred = math_ops.matmul(embedding_ops.embedding_lookup([var0], [0]), x)
+        return pred * pred
+
+      # pylint: enable=cell-var-from-loop
+
+      opt = momentum_lib.MomentumOptimizer(learning_rate=1.0, momentum=0.0)
+      sgd_op = opt.minimize(loss)
+      self.evaluate(variables.global_variables_initializer())
+      # Run 1 step of sgd
+      self.evaluate(sgd_op)
+      # Validate updated params
+      self.assertAllCloseAccordingToType([[-111, -138]], self.evaluate(var0))
+
+  @test_util.run_in_graph_and_eager_modes
+  def testMinimizeWith2DIndicesForEmbeddingLookup(self):
+    # This test invokes the ResourceSparseApplyMomentum operation, which
+    # did not have a registered GPU kernel as of April 2018. With graph
+    # execution, the placement algorithm notices this and automatically
+    # places the variable in CPU (host) memory. With eager execution,
+    # the variable would be placed in GPU memory if available, which
+    # would then conflict with the future invocation of the
+    # ResourceSparseApplyMomentum operation.
+    # To work around this discrepancy, for now we force the variable
+    # to be placed on CPU.
+    with ops.device("/cpu:0"):
+      var0 = resource_variable_ops.ResourceVariable(array_ops.ones([2, 2]))
+
+    def loss():
+      return math_ops.reduce_sum(embedding_ops.embedding_lookup(var0, [[1]]))
+
+    opt = momentum_lib.MomentumOptimizer(learning_rate=1.0, momentum=0.0)
+    sgd_op = opt.minimize(loss)
+    self.evaluate(variables.global_variables_initializer())
+    self.evaluate(sgd_op)
+    self.assertAllCloseAccordingToType([[1, 1], [0, 0]], self.evaluate(var0))
+
+  def testTensorLearningRateAndMomentum(self):
+    for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+      # train.MomentumOptimizer is V1 only API.
+      with ops.Graph().as_default(), self.cached_session():
+        var0 = variables.Variable([1.0, 2.0], dtype=dtype)
+        var1 = variables.Variable([3.0, 4.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.1], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.01], dtype=dtype)
+        mom_opt = momentum_lib.MomentumOptimizer(
+            learning_rate=constant_op.constant(2.0),
+            momentum=constant_op.constant(0.9))
+        mom_update = mom_opt.apply_gradients(
+            zip([grads0, grads1], [var0, var1]))
+        self.evaluate(variables.global_variables_initializer())
+        # Check we have slots
+        self.assertEqual(["momentum"], mom_opt.get_slot_names())
+        slot0 = mom_opt.get_slot(var0, "momentum")
+        self.assertEqual(slot0.get_shape(), var0.get_shape())
+        self.assertFalse(slot0 in variables.trainable_variables())
+        slot1 = mom_opt.get_slot(var1, "momentum")
+        self.assertEqual(slot1.get_shape(), var1.get_shape())
+        self.assertFalse(slot1 in variables.trainable_variables())
+
+        # 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
+        mom_update.run()
+        # Check that the momentum accumulators have been updated.
+        self.assertAllCloseAccordingToType(
+            np.array([0.1, 0.1]), self.evaluate(slot0))
+        self.assertAllCloseAccordingToType(
+            np.array([0.01, 0.01]), 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.
+        mom_update.run()
+        # Check that the momentum accumulators have been updated.
+        self.assertAllCloseAccordingToType(
+            np.array([(0.9 * 0.1 + 0.1), (0.9 * 0.1 + 0.1)]),
+            self.evaluate(slot0))
+        self.assertAllCloseAccordingToType(
+            np.array([(0.9 * 0.01 + 0.01), (0.9 * 0.01 + 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 _dbParamsMom01(self):
+    """Return dist-belief momentum values.
+
+    Return values been generated from the dist-belief momentum unittest,
+    running with a learning rate of 0.1 and a momentum of 0.1.
+
+    These values record how a parameter vector of size 10, initialized with 0.0,
+    gets updated with 10 consecutive momentum steps.  It uses random gradients.
+
+    Returns:
+      db_grad: The gradients to apply
+      db_out: The parameters after the momentum update.
+    """
+    db_grad = [[]] * 10
+    db_out = [[]] * 10
+    # pylint: disable=line-too-long
+    db_grad[0] = [
+        0.00096264342, 0.17914793, 0.93945462, 0.41396621, 0.53037018,
+        0.93197989, 0.78648776, 0.50036013, 0.55345792, 0.96722615
+    ]
+    db_out[0] = [
+        -9.6264346e-05, -0.017914793, -0.093945466, -0.041396622, -0.053037018,
+        -0.093197994, -0.078648776, -0.050036013, -0.055345792, -0.096722618
+    ]
+    db_grad[1] = [
+        0.17075552, 0.88821375, 0.20873757, 0.25236958, 0.57578111, 0.15312378,
+        0.5513742, 0.94687688, 0.16012503, 0.22159521
+    ]
+    db_out[1] = [
+        -0.017181443, -0.10852765, -0.12421377, -0.070773244, -0.11591884,
+        -0.11783017, -0.14165108, -0.14972731, -0.076892875, -0.1285544
+    ]
+    db_grad[2] = [
+        0.35077485, 0.47304362, 0.44412705, 0.44368884, 0.078527533, 0.81223965,
+        0.31168157, 0.43203235, 0.16792089, 0.24644311
+    ]
+    db_out[2] = [
+        -0.053967446, -0.1648933, -0.1716533, -0.1180798, -0.13005978,
+        -0.20151734, -0.17911947, -0.20289968, -0.095839672, -0.15638189
+    ]
+    db_grad[3] = [
+        0.9694621, 0.75035888, 0.28171822, 0.83813518, 0.53807181, 0.3728098,
+        0.81454384, 0.03848977, 0.89759839, 0.93665648
+    ]
+    db_out[3] = [
+        -0.15459226, -0.24556576, -0.20456907, -0.20662397, -0.18528105,
+        -0.24716705, -0.2643207, -0.21206589, -0.18749419, -0.2528303
+    ]
+    db_grad[4] = [
+        0.38578293, 0.8536852, 0.88722926, 0.66276771, 0.13678469, 0.94036359,
+        0.69107032, 0.81897682, 0.5433259, 0.67860287
+    ]
+    db_out[4] = [
+        -0.20323303, -0.33900154, -0.29658359, -0.28175515, -0.20448165,
+        -0.34576839, -0.34194785, -0.29488021, -0.25099224, -0.33033544
+    ]
+    db_grad[5] = [
+        0.27885768, 0.76100707, 0.24625534, 0.81354135, 0.18959245, 0.48038563,
+        0.84163809, 0.41172323, 0.83259648, 0.44941229
+    ]
+    db_out[5] = [
+        -0.23598288, -0.42444581, -0.33041057, -0.3706224, -0.22536094,
+        -0.40366709, -0.43387437, -0.34433398, -0.34060168, -0.38302717
+    ]
+    db_grad[6] = [
+        0.27233034, 0.056316052, 0.5039115, 0.24105175, 0.35697976, 0.75913221,
+        0.73577434, 0.16014607, 0.57500273, 0.071136251
+    ]
+    db_out[6] = [
+        -0.26649091, -0.43862185, -0.38418442, -0.40361428, -0.26314685,
+        -0.48537019, -0.51664448, -0.36529395, -0.40706289, -0.39540997
+    ]
+    db_grad[7] = [
+        0.58697265, 0.2494842, 0.08106143, 0.39954534, 0.15892942, 0.12683646,
+        0.74053431, 0.16033, 0.66625422, 0.73515922
+    ]
+    db_out[7] = [
+        -0.32823896, -0.46498787, -0.39766794, -0.446868, -0.28281838,
+        -0.50622416, -0.59897494, -0.38342294, -0.48033443, -0.47016418
+    ]
+    db_grad[8] = [
+        0.8215279, 0.41994119, 0.95172721, 0.68000203, 0.79439718, 0.43384039,
+        0.55561525, 0.22567581, 0.93331909, 0.29438227
+    ]
+    db_out[8] = [
+        -0.41656655, -0.50961858, -0.49418902, -0.51919359, -0.36422527,
+        -0.55169362, -0.6627695, -0.40780342, -0.58099347, -0.50707781
+    ]
+    db_grad[9] = [
+        0.68297005, 0.67758518, 0.1748755, 0.13266537, 0.70697063, 0.055731893,
+        0.68593478, 0.50580865, 0.12602448, 0.093537711
+    ]
+    db_out[9] = [
+        -0.49369633, -0.58184016, -0.52132869, -0.5396927, -0.44306302,
+        -0.56181377, -0.73774242, -0.46082234, -0.60366184, -0.52012295
+    ]
+    # pylint: enable=line-too-long
+    return db_grad, db_out
+
+  def testLikeDistBeliefMom01(self):
+    # train.MomentumOptimizer is V1 only API.
+    with ops.Graph().as_default(), self.cached_session():
+      db_grad, db_out = self._dbParamsMom01()
+      num_samples = len(db_grad)
+      var0 = variables.Variable([0.0] * num_samples)
+      grads0 = constant_op.constant([0.0] * num_samples)
+      mom_opt = momentum_lib.MomentumOptimizer(learning_rate=0.1, momentum=0.1)
+      mom_update = mom_opt.apply_gradients(zip([grads0], [var0]))
+      self.evaluate(variables.global_variables_initializer())
+      for i in range(num_samples):
+        mom_update.run(feed_dict={grads0: db_grad[i]})
+        self.assertAllClose(np.array(db_out[i]), self.evaluate(var0))
+
+  def testSparse(self):
+    for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+      # train.MomentumOptimizer is V1 only API.
+      with ops.Graph().as_default(), self.cached_session():
+        var0 = variables.Variable(array_ops.zeros([4, 2], dtype=dtype))
+        var1 = variables.Variable(constant_op.constant(1.0, dtype, [4, 2]))
+        grads0 = indexed_slices.IndexedSlices(
+            constant_op.constant([[.1, .1]], dtype=dtype),
+            constant_op.constant([1]), constant_op.constant([4, 2]))
+        grads1 = indexed_slices.IndexedSlices(
+            constant_op.constant([[.01, .01], [.01, .01]], dtype=dtype),
+            constant_op.constant([2, 3]), constant_op.constant([4, 2]))
+        mom_opt = momentum_lib.MomentumOptimizer(
+            learning_rate=2.0, momentum=0.9)
+        mom_update = mom_opt.apply_gradients(
+            zip([grads0, grads1], [var0, var1]))
+        self.evaluate(variables.global_variables_initializer())
+
+        # Check we have slots
+        self.assertEqual(["momentum"], mom_opt.get_slot_names())
+        slot0 = mom_opt.get_slot(var0, "momentum")
+        self.assertEqual(slot0.get_shape(), var0.get_shape())
+        slot1 = mom_opt.get_slot(var1, "momentum")
+        self.assertEqual(slot1.get_shape(), var1.get_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
+        mom_update.run()
+        # Check that the momentum accumulators have been updated.
+        self.assertAllCloseAccordingToType(
+            np.array([0, 0]),
+            self.evaluate(slot0)[0])
+        self.assertAllCloseAccordingToType(
+            np.array([.1, .1]),
+            self.evaluate(slot0)[1])
+        self.assertAllCloseAccordingToType(
+            np.array([.01, .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.
+        mom_update.run()
+        # 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.1 + 0.1), (0.9 * 0.1 + 0.1)]),
+            self.evaluate(slot0)[1])
+        self.assertAllCloseAccordingToType(
+            np.array([(0.9 * 0.01 + 0.01), (0.9 * 0.01 + 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):
+    for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+      # train.MomentumOptimizer is V1 only API.
+      with ops.Graph().as_default(), self.cached_session():
+        var0 = variables.Variable([1.0, 2.0], dtype=dtype)
+        var1 = variables.Variable([3.0, 4.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.1], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.01], dtype=dtype)
+        mom_opt = momentum_lib.MomentumOptimizer(
+            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(variables.global_variables_initializer())
+
+        self.assertEqual(["momentum"], mom_opt.get_slot_names())
+        slot0 = mom_opt.get_slot(var0, "momentum")
+        self.assertEqual(slot0.get_shape(), var0.get_shape())
+        slot1 = mom_opt.get_slot(var1, "momentum")
+        self.assertEqual(slot1.get_shape(), var1.get_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
+        mom_update1.run()
+        # Check that the momentum accumulators have been updated.
+        self.assertAllCloseAccordingToType(
+            np.array([0.1, 0.1]), self.evaluate(slot0))
+        self.assertAllCloseAccordingToType(
+            np.array([0.01, 0.01]), 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.
+        mom_update2.run()
+        # Check that the momentum accumulators have been updated.
+        self.assertAllCloseAccordingToType(
+            np.array([(0.9 * 0.1 + 0.1), (0.9 * 0.1 + 0.1)]),
+            self.evaluate(slot0))
+        self.assertAllCloseAccordingToType(
+            np.array([(0.9 * 0.01 + 0.01), (0.9 * 0.01 + 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))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/moving_averages.py

@@ -0,0 +1,689 @@
+# 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.
+# ==============================================================================
+"""Maintain moving averages of parameters."""
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import reduce_util as ds_reduce_util
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import state_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.ops import variables
+from tensorflow.python.training import slot_creator
+from tensorflow.python.util.tf_export import tf_export
+from tensorflow.tools.docs import doc_controls
+
+
+@tf_export("__internal__.train.assign_moving_average", v1=[])
+def assign_moving_average(variable, value, decay, zero_debias=True, name=None):
+  """Compute the moving average of a variable.
+
+  The moving average of 'variable' updated with 'value' is:
+    variable * decay + value * (1 - decay)
+
+  The returned Operation sets 'variable' to the newly computed moving average,
+  by performing this subtraction:
+     variable -= (1 - decay) * (variable - value)
+
+  Since variables that are initialized to a `0` value will be `0` biased,
+  `zero_debias` optionally enables scaling by the mathematically correct
+  debiasing factor of
+    1 - decay ** num_updates
+  See Section 3 of (Kingma et al., 2015) for more details.
+
+  The names of the debias shadow variables, by default, include both the scope
+  they were created in and the scope of the variables they debias. They are also
+  given a uniquifying-suffix.
+
+  E.g.:
+
+  ```
+    with tf.compat.v1.variable_scope('scope1'):
+      with tf.compat.v1.variable_scope('scope2'):
+        var = tf.compat.v1.get_variable('foo')
+        update_1 = tf.assign_moving_average(var, 0.0, 1.0)
+        update_2 = tf.assign_moving_average(var, 0.0, 0.9)
+
+    # var.name: 'scope1/scope2/foo'
+    # shadow var names: 'scope1/scope2/scope1/scope2/foo/biased'
+    #                   'scope1/scope2/scope1/scope2/foo/biased_1'
+  ```
+
+  Args:
+    variable: A Variable.
+    value: A tensor with the same shape as 'variable'.
+    decay: A float `Tensor` or float value. The moving average decay.
+    zero_debias: A python bool. If true, assume the variable is 0-initialized
+      and unbias it, as in (Kingma et al., 2015). See docstring in
+        `_zero_debias` for more details.
+    name: Optional name of the returned operation.
+
+  Returns:
+    A tensor which if evaluated will compute and return the new moving average.
+
+  References:
+    Adam - A Method for Stochastic Optimization:
+      [Kingma et al., 2015](https://arxiv.org/abs/1412.6980)
+      ([pdf](https://arxiv.org/pdf/1412.6980.pdf))
+  """
+  with ops.name_scope(name, "AssignMovingAvg",
+                      [variable, value, decay]) as scope:
+    decay = ops.convert_to_tensor(1.0 - decay, name="decay")
+    if decay.dtype != variable.dtype.base_dtype:
+      decay = math_ops.cast(decay, variable.dtype.base_dtype)
+
+    def update_fn(v, value):
+      return state_ops.assign_sub(v, (v - value) * decay, name=scope)
+
+    def update(strategy, v, value):
+      if zero_debias:
+        return _zero_debias(strategy, v, value, decay)
+      else:
+        return _update(strategy, v, update_fn, args=(value,))
+
+    replica_context = distribute_lib.get_replica_context()
+    if replica_context:
+      # In a replica context, we update variable using the mean of value across
+      # replicas.
+      def merge_fn(strategy, v, value):
+        value = strategy.extended.reduce_to(ds_reduce_util.ReduceOp.MEAN, value,
+                                            v)
+        return update(strategy, v, value)
+
+      return replica_context.merge_call(merge_fn, args=(variable, value))
+    else:
+      strategy = distribute_lib.get_cross_replica_context()
+      return update(strategy, variable, value)
+
+
+def weighted_moving_average(value,
+                            decay,
+                            weight,
+                            truediv=True,
+                            collections=None,
+                            name=None):
+  """Compute the weighted moving average of `value`.
+
+  Conceptually, the weighted moving average is:
+    `moving_average(value * weight) / moving_average(weight)`,
+  where a moving average updates by the rule
+    `new_value = decay * old_value + (1 - decay) * update`
+  Internally, this Op keeps moving average variables of both `value * weight`
+  and `weight`.
+
+  Args:
+    value: A numeric `Tensor`.
+    decay: A float `Tensor` or float value. The moving average decay.
+    weight:  `Tensor` that keeps the current value of a weight. Shape should be
+      able to multiply `value`.
+    truediv:  Boolean, if `True`, dividing by `moving_average(weight)` is
+      floating point division.  If `False`, use division implied by dtypes.
+    collections:  List of graph collections keys to add the internal variables
+      `value * weight` and `weight` to. Defaults to
+      `[GraphKeys.GLOBAL_VARIABLES]`.
+    name: Optional name of the returned operation. Defaults to
+      "WeightedMovingAvg".
+
+  Returns:
+    An Operation that updates and returns the weighted moving average.
+  """
+  # Unlike assign_moving_average, the weighted moving average doesn't modify
+  # user-visible variables. It is the ratio of two internal variables, which are
+  # moving averages of the updates.  Thus, the signature of this function is
+  # quite different than assign_moving_average.
+  if collections is None:
+    collections = [ops.GraphKeys.GLOBAL_VARIABLES]
+  with variable_scope.variable_scope(name, "WeightedMovingAvg",
+                                     [value, weight, decay]) as scope:
+    value_x_weight_var = variable_scope.get_variable(
+        "value_x_weight",
+        shape=value.get_shape(),
+        dtype=value.dtype,
+        initializer=init_ops.zeros_initializer(),
+        trainable=False,
+        collections=collections)
+    weight_var = variable_scope.get_variable(
+        "weight",
+        shape=weight.get_shape(),
+        dtype=weight.dtype,
+        initializer=init_ops.zeros_initializer(),
+        trainable=False,
+        collections=collections)
+    numerator = assign_moving_average(
+        value_x_weight_var, value * weight, decay, zero_debias=False)
+    denominator = assign_moving_average(
+        weight_var, weight, decay, zero_debias=False)
+
+    if truediv:
+      return math_ops.truediv(numerator, denominator, name=scope.name)
+    else:
+      return math_ops.divide(numerator, denominator, name=scope.name)
+
+
+def _update(strategy, var, update_fn, args):
+  """Applies updates depending on the context."""
+  assert distribute_lib.in_cross_replica_context(), (
+      "_update can only be called in cross-replica context")
+  if distribute_lib.get_update_replica_id() is not None:
+    # Call update_fn on var to delegate the implementation. We expect `var` will
+    # do the right thing in update context, e.g, if `var` is a MirroredVariable,
+    # it should pick its component variable based on `update_replica_id` and
+    # only update that.
+    return update_fn(var, *args)
+  else:
+    return strategy.extended.update(var, update_fn, args)
+
+
+def _zero_debias(strategy, unbiased_var, value, decay):
+  """Compute the delta required for a debiased Variable.
+
+  All exponential moving averages initialized with Tensors are initialized to 0,
+  and therefore are biased to 0. Variables initialized to 0 and used as EMAs are
+  similarly biased. This function creates the debias updated amount according to
+  a scale factor, as in (Kingma et al., 2015).
+
+  To demonstrate the bias the results from 0-initialization, take an EMA that
+  was initialized to `0` with decay `b`. After `t` timesteps of seeing the
+  constant `c`, the variable have the following value:
+
+  ```
+    EMA = 0*b^(t) + c*(1 - b)*b^(t-1) + c*(1 - b)*b^(t-2) + ...
+        = c*(1 - b^t)
+  ```
+
+  To have the true value `c`, we would divide by the scale factor `1 - b^t`.
+
+  In order to perform debiasing, we use two shadow variables. One keeps track of
+  the biased estimate, and the other keeps track of the number of updates that
+  have occurred.
+
+  Args:
+    strategy: `Strategy` used to create and update variables.
+    unbiased_var: A Variable representing the current value of the unbiased EMA.
+    value: A Tensor representing the most recent value.
+    decay: A Tensor representing `1-decay` for the EMA.
+
+  Returns:
+    The amount that the unbiased variable should be updated. Computing this
+    tensor will also update the shadow variables appropriately.
+
+  References:
+    Adam - A Method for Stochastic Optimization:
+      [Kingma et al., 2015](https://arxiv.org/abs/1412.6980)
+      ([pdf](https://arxiv.org/pdf/1412.6980.pdf))
+
+  """
+  with variable_scope.variable_scope(
+      unbiased_var.name[:-len(":0")], values=[unbiased_var, value, decay]):
+    with ops.init_scope():
+      biased_initializer = init_ops.zeros_initializer()
+      local_step_initializer = init_ops.zeros_initializer()
+
+    def _maybe_get_unique(name):
+      """Get name for a unique variable, if not `reuse=True`."""
+      if variable_scope.get_variable_scope().reuse:
+        return name
+      vs_vars = [
+          x.op.name
+          for x in variable_scope.get_variable_scope().global_variables()
+      ]
+      full_name = variable_scope.get_variable_scope().name + "/" + name
+      if full_name not in vs_vars:
+        return name
+      idx = 1
+      while full_name + ("_%d" % idx) in vs_vars:
+        idx += 1
+      return name + ("_%d" % idx)
+
+    with strategy.extended.colocate_vars_with(unbiased_var):
+      biased_var = variable_scope.get_variable(
+          _maybe_get_unique("biased"),
+          initializer=biased_initializer,
+          shape=unbiased_var.get_shape(),
+          dtype=unbiased_var.dtype,
+          trainable=False)
+      local_step = variable_scope.get_variable(
+          _maybe_get_unique("local_step"),
+          shape=[],
+          dtype=unbiased_var.dtype,
+          initializer=local_step_initializer,
+          trainable=False)
+
+  def update_fn(v, value, biased_var, local_step):
+    update_biased = state_ops.assign_sub(biased_var,
+                                         (biased_var - value) * decay)
+    update_local_step = local_step.assign_add(1)
+
+    # This function gets `1 - decay`, so use `1.0 - decay` in the exponent.
+    bias_factor = 1 - math_ops.pow(1.0 - decay, update_local_step)
+    return state_ops.assign(
+        v, update_biased / bias_factor, name=ops.get_name_scope() + "/")
+
+  return _update(
+      strategy, unbiased_var, update_fn, args=(value, biased_var, local_step))
+
+
+@tf_export("train.ExponentialMovingAverage")
+class ExponentialMovingAverage:
+  """Maintains moving averages of variables by employing an exponential decay.
+
+  When training a model, it is often beneficial to maintain moving averages of
+  the trained parameters.  Evaluations that use averaged parameters sometimes
+  produce significantly better results than the final trained values.
+
+  The `apply()` method adds shadow copies of trained variables the first time
+  it is called, and maintains a moving average of the trained variables in
+  their shadow copies at every additional invocation.
+  It should generally be called immediately after creating the model weights,
+  and then after each training step.
+
+  The `average()` method gives access to the shadow variables.
+  It allows you to use the moving averages in place of the last trained values
+  for evaluations, by loading the moving averages into your model via
+  `var.assign(ema.average(var))`.
+  Additionally, although `ExponentialMovingAverage`
+  objects are not directly trackable by checkpoints,
+  `average()` returns the moving average variables for your model weights,
+  which you can then checkpoint. (There is an example
+  of this near the bottom of this docstring).
+  So, `average()` is useful when
+  building an evaluation model, or when restoring a model from a checkpoint
+  file.
+
+  The moving averages are computed using exponential decay.  You specify the
+  decay value (as a scalar float value, `Tensor`, or `Variable`) when creating
+  the `ExponentialMovingAverage` object.  The shadow variables are initialized
+  with the same initial values as the trained variables.  When you run `apply`
+  to update the moving averages, each shadow variable is updated with the
+  formula:
+
+    `shadow_variable -= (1 - decay) * (shadow_variable - variable)`
+
+  This is mathematically equivalent to the classic formula below, but the use
+  of an `assign_sub` op (the `"-="` in the formula) allows concurrent lockless
+  updates to the variables:
+
+    `shadow_variable = decay * shadow_variable + (1 - decay) * variable`
+
+  Reasonable values for `decay` are close to 1.0, typically in the
+  multiple-nines range: 0.999, 0.9999, etc.
+
+  To have fine-grained control over the value of the decay parameter during
+  training, pass a scalar `tf.Variable` as the `decay` value to the constructor,
+  and update the variable as needed.
+
+  Example usage when creating a training model:
+
+  ```python
+  # Create variables.
+  var0 = tf.Variable(...)
+  var1 = tf.Variable(...)
+  # ... use the variables to build a training model...
+
+  # Create an ExponentialMovingAverage object
+  ema = tf.train.ExponentialMovingAverage(decay=0.9999)
+
+  # The first `apply` creates the shadow variables that hold the moving averages
+  ema.apply([var0, var1])
+
+  # grab the moving averages for checkpointing purposes or to be able to
+  # load the moving averages into the model weights
+  averages = [ema.average(var0), ema.average(var1)]
+
+  ...
+  def train_step(...):
+  ...
+    # Apply the optimizer.
+    opt.minimize(my_loss, [var0, var1])
+
+    # Update the moving averages
+    # of var0 and var1 with additional calls to `apply`
+    ema.apply([var0, var1])
+
+  ...train the model by running train_step multiple times...
+  ```
+
+  There are several ways to use the moving averages for evaluations:
+
+  1. Assign the values of the shadow variables to your model variables with
+     `Variable.assign(...)` before evaluating your
+     model. You can use the `average()`
+     method to get the shadow variable for a given variable. To continue
+     training after using this approach, make sure to record the unaveraged
+     weights and restore them before continuing to train. You can see the
+     tensorflow-addons' MovingAverage optimizer's `swap_weights` method for
+     one example of how to swap variables efficiently in distributed settings:
+     https://github.com/tensorflow/addons/blob/v0.13.0/tensorflow_addons/optimizers/moving_average.py#L151
+  2. Make sure to checkpoint out your moving average variables in your
+     `tf.train.Checkpoint`. At evaluation time, create your shadow variables and
+     use `tf.train.Checkpoint` to restore the moving averages into the shadow
+     variables. Then, load the moving averages into the actual model weights via
+     `var.assign(moving_avg)`.
+  3. Checkpoint out your moving average variables in your `tf.train.Checkpoint`.
+     For evaluation, restore your model weights directly from the moving
+     averages instead of from the non-averaged weights.
+     Caution: If you choose this approach, include only the object-graph paths
+     to the averaged path in your checkpoint restore.
+     If you point both the unaveraged and averaged paths in a checkpoint
+     restore to the same variables, it is hard to reason about whether your
+     model will restore the averaged or non-averaged variables.
+
+  Example of saving out then restoring the shadow variable values:
+
+  ```python
+  # Create variables.
+  var0 = tf.Variable(...)
+  var1 = tf.Variable(...)
+  # ... use the variables to build a training model...
+
+  # Create an ExponentialMovingAverage object, create the shadow variables,
+  # and grab the moving averages for checkpointing purposes.
+  # (The ExponentialMovingAverage object itself is not checkpointable)
+  ema = tf.train.ExponentialMovingAverage(decay=0.9999)
+  ema.apply([var0, var1])
+  avg_var0 = ema.average(var0)
+  avg_var1 = ema.average(var1)
+
+  # Create a Checkpoint that will manage the model weights and the averages,
+  checkpoint = tf.train.Checkpoint(model_weights=[var0, var1],
+                                   averaged_weights=[avg_var0, avg_var1])
+  ... # Do training
+
+  # Save out the checkpoint including the model weights and the moving averages
+  checkpoint.save(...)
+  ```
+
+  Restore option: restore all averaged & non-averaged weights, then load
+  moving averages into the model via `var.assign()`
+  ```python
+  # Create variables.
+  var0 = tf.Variable(...)
+  var1 = tf.Variable(...)
+  # ... use the variables to build a training model...
+
+  # Create an ExponentialMovingAverage object, create the shadow variables,
+  # and grab the moving averages for checkpoint restore purposes.
+  # (The ExponentialMovingAverage object itself is not checkpointable)
+  ema = tf.train.ExponentialMovingAverage(decay=0.9999)
+  ema.apply([var0, var1])
+  avg_var0 = ema.average(var0)
+  avg_var1 = ema.average(var1)
+
+  # Create a Checkpoint that will manage the model weights and the averages,
+  checkpoint = tf.train.Checkpoint(model_weights=[var0, var1],
+                                   averaged_weights=[avg_var0, avg_var1])
+  checkpoint.restore(...)
+  var0.assign(avg_var0)
+  var1.assign(avg_var1)
+  # var0 and var1 now hold the moving average values
+  ```
+
+  Restore option: Directly restore the moving averages into the model weights.
+  ```python
+  # Create variables.
+  var0 = tf.Variable(...)
+  var1 = tf.Variable(...)
+  # ... use the variables to build a training model...
+
+  # Create a Checkpoint that will manage two objects with trackable state,
+  checkpoint = tf.train.Checkpoint(averaged_weights=[var0, var1])
+  checkpoint.restore(...)
+  # var0 and var1 now hold the moving average values
+  ```
+  """
+
+  def __init__(self,
+               decay,
+               num_updates=None,
+               zero_debias=False,
+               name="ExponentialMovingAverage"):
+    """Creates a new ExponentialMovingAverage object.
+
+    The `apply()` method has to be called to create shadow variables.
+    Follow-on calls to the `apply()` method will update the moving averages
+    in the shadow variables.
+    (In TF 1.x graphs `apply()` will return an update op to update
+    the moving averages which must be explicitly run).
+
+    The optional `num_updates` parameter allows one to tweak the decay rate
+    dynamically. It is typical to pass the count of training steps, usually
+    kept in a variable that is incremented at each step, in which case the
+    decay rate is lower at the start of training.  This makes moving averages
+    move faster.  If passed, the actual decay rate used is:
+
+      `min(decay, (1 + num_updates) / (10 + num_updates))`
+
+    Args:
+      decay: A scalar float value, `Tensor`, or `Variable`. The decay parameter.
+      num_updates: Optional count of number of updates applied to variables.
+      zero_debias: If `True`, zero debias moving-averages that are initialized
+        with tensors. (Note: moving averages may not be initialized with
+        non-variable tensors when eager execution is enabled).
+      name: String. Optional prefix name to use for the name of ops added in
+        `apply()`.
+    """
+    self._decay = decay
+    self._num_updates = num_updates
+    self._zero_debias = zero_debias
+    self._name = name
+    self._averages = {}
+
+  @property
+  def name(self):
+    """The name of this ExponentialMovingAverage object."""
+    return self._name
+
+  def apply(self, var_list=None):
+    """Maintains moving averages of variables.
+
+    `var_list` must be a list of `Variable` objects.  This method
+    creates shadow variables (holding the moving averages)
+    for all elements of `var_list`, and
+    updates the moving averages using the current `var_list` values. Shadow
+    variables for `Variable` objects are initialized to the variable's initial
+    value.
+
+    Shadow variables are created with `trainable=False`. To access them you
+    can use the EMA object's `average` method. Note that `EMA` objects are
+    not trackable by checkpoints, so if you want to checkpoint or restore the
+    moving variables you will need to manually grab the shadow
+    variables via `average()` and assign them as `tf.Module` properties or
+    directly pass them to your `tf.train.Checkpoint`.
+
+    Note that `apply()` can be called multiple times. When eager execution is
+    enabled each call to apply will update the variables once, so this needs to
+    be called in a loop.
+
+    In legacy TF 1.x graphs, this method returns an op that updates all
+    shadow variables from the current value of their associated variables. In
+    TF 1.x graphs without automatically control dependencies this op needs to be
+    manually run.
+
+    Args:
+      var_list: A list of Variable objects. The variables
+        must be of types bfloat16, float16, float32, or float64.
+        (In legacy TF 1.x graphs these may be tensors, but this is unsupported
+        when eager execution is enabled.)
+
+    Returns:
+      An Operation that updates the moving averages.
+
+    Raises:
+      TypeError: If the arguments are not an allowed type.
+    """
+    # TODO(touts): op_scope
+    if var_list is None:
+      var_list = variables.trainable_variables()
+    for v in var_list:
+      if (isinstance(v, tensor.Tensor)
+          and ops.executing_eagerly_outside_functions()):
+        raise TypeError(
+            "tf.train.ExponentialMovingAverage does not support non-Variable"
+            " tensors when eager execution is enabled.")
+    zero_debias_true = set()  # set of vars to set `zero_debias=True`
+    for var in var_list:
+      if var.dtype.base_dtype not in [
+          dtypes.bfloat16, dtypes.float16, dtypes.float32, dtypes.float64
+      ]:
+        raise TypeError("The variables must be half, float, or double: %s" %
+                        var.name)
+
+      if var.ref() not in self._averages:
+        # For variables: to lower communication bandwidth across devices we keep
+        # the moving averages on the same device as the variables. For other
+        # tensors, we rely on the existing device allocation mechanism.
+        with ops.init_scope():
+          if isinstance(var, variables.Variable):
+            with ops.device(var.device):
+              initialized_value = cond.cond(
+                  variable_v1.is_variable_initialized(var), var.read_value,
+                  lambda: var.initial_value)  # pylint: disable=cell-var-from-loop
+            avg = slot_creator.create_slot(
+                var,
+                initialized_value,
+                self.name,
+                colocate_with_primary=True,
+                copy_xla_sharding=True)
+            # NOTE(mrry): We only add `tf.Variable` objects to the
+            # `MOVING_AVERAGE_VARIABLES` collection.
+            ops.add_to_collection(ops.GraphKeys.MOVING_AVERAGE_VARIABLES, var)
+          else:
+            avg = slot_creator.create_zeros_slot(
+                var,
+                self.name,
+                colocate_with_primary=(var.op.type in [
+                    "Variable", "VariableV2", "VarHandleOp"
+                ]),
+                copy_xla_sharding=True)
+            if self._zero_debias:
+              zero_debias_true.add(avg.ref())
+        self._averages[var.ref()] = avg
+
+    with ops.name_scope(self.name) as scope:
+      decay = ops.convert_to_tensor(
+          self._decay, dtype=dtypes.float32, name="decay")
+      if self._num_updates is not None:
+        num_updates = math_ops.cast(
+            self._num_updates, dtypes.float32, name="num_updates")
+        decay = math_ops.minimum(decay,
+                                 (1.0 + num_updates) / (10.0 + num_updates))
+      updates = []
+      for var in var_list:
+        avg = self._averages[var.ref()]
+        zero_debias = avg.ref() in zero_debias_true
+        updates.append(assign_moving_average(avg, var, decay, zero_debias))
+      return control_flow_ops.group(*updates, name=scope)
+
+  def average(self, var):
+    """Returns the `Variable` holding the average of `var`.
+
+    Args:
+      var: A `Variable` object.
+
+    Returns:
+      A `Variable` object or `None` if the moving average of `var`
+      is not maintained.
+    """
+    return self._averages.get(var.ref(), None)
+
+  @doc_controls.do_not_generate_docs
+  def average_name(self, var):
+    """[Meant for TF1] Returns name of `Variable` holding the average for `var`.
+
+    (Designed to work with legacy `tf.compat.v1.train.Saver`, it is sensitive to
+    specific variable names and not recommended for TF2)
+
+    The typical scenario for `ExponentialMovingAverage` is to compute moving
+    averages of variables during training, and restore the variables from the
+    computed moving averages during evaluations.
+
+    To restore variables, you have to know the name of the shadow variables.
+    That name and the original variable can then be passed to a `Saver()` object
+    to restore the variable from the moving average value with:
+      `saver = tf.compat.v1.train.Saver({ema.average_name(var): var})`
+
+    `average_name()` can be called whether or not `apply()` has been called.
+
+    Args:
+      var: A `Variable` object.
+
+    Returns:
+      A string: The name of the variable that will be used or was used
+      by the `ExponentialMovingAverage class` to hold the moving average of
+      `var`.
+    """
+    if var.ref() in self._averages:
+      return self._averages[var.ref()].name[:-len(":0")]
+    return ops.get_default_graph().unique_name(
+        var.name[:-len(":0")] + "/" + self.name, mark_as_used=False)
+
+  @doc_controls.do_not_generate_docs
+  def variables_to_restore(self, moving_avg_variables=None):
+    """[Designed for TF 1.x] Returns a map of names to `Variables` to restore.
+
+    (Designed to work with legacy `tf.compat.v1.train.Saver`, sensitive to
+    specific variable names and not recommended for TF2)
+
+    If a variable has a moving average, use the moving average variable name as
+    the restore name; otherwise, use the variable name.
+
+    For example,
+
+    ```python
+      variables_to_restore = ema.variables_to_restore()
+      saver = tf.compat.v1.train.Saver(variables_to_restore)
+    ```
+
+    Below is an example of such mapping:
+
+    ```
+      conv/batchnorm/gamma/ExponentialMovingAverage: conv/batchnorm/gamma,
+      conv_4/conv2d_params/ExponentialMovingAverage: conv_4/conv2d_params,
+      global_step: global_step
+    ```
+
+    Args:
+      moving_avg_variables: a list of variables that require to use of the
+        moving average variable name to be restored. If None, it will default to
+        variables.moving_average_variables() + variables.trainable_variables()
+
+    Returns:
+      A map from restore_names to variables. The restore_name is either the
+      original or the moving average version of the variable name, depending
+      on whether the variable name is in the `moving_avg_variables`.
+    """
+    name_map = {}
+    if moving_avg_variables is None:
+      # Include trainable variables and variables which have been explicitly
+      # added to the moving_average_variables collection.
+      moving_avg_variables = variables.trainable_variables()
+      moving_avg_variables += variables.moving_average_variables()
+    # Remove duplicates
+    moving_avg_variables = set(v.ref() for v in moving_avg_variables)
+    # Collect all the variables with moving average,
+    for v in moving_avg_variables:
+      name_map[self.average_name(v.deref())] = v.deref()
+    # Make sure we restore variables without moving averages as well.
+    moving_avg_variable_names = set(
+        v.deref().name for v in moving_avg_variables)
+    for v in list(set(variables.global_variables())):
+      if v.name not in moving_avg_variable_names and v.op.name not in name_map:
+        name_map[v.op.name] = v
+    return name_map

tensorflow/python/training/moving_averages_test.py

@@ -0,0 +1,593 @@
+# 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 moving_averages.py."""
+
+import numpy as np
+
+from tensorflow.python.compiler.xla.experimental import xla_sharding
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_state_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import moving_averages
+from tensorflow.python.training import saver as saver_lib
+
+
+class MovingAveragesTest(test.TestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def testAssignMovingAverageWithoutZeroDebias(self):
+    var = variables.Variable([10.0, 11.0])
+    val = constant_op.constant([1.0, 2.0], dtypes.float32)
+    decay = 0.25
+    if context.executing_eagerly():
+      self.assertAllClose([10.0, 11.0], self.evaluate(var))
+      assign = moving_averages.assign_moving_average(
+          var, val, decay, zero_debias=False)
+      self.assertAllClose(
+          [10.0 * 0.25 + 1.0 * (1.0 - 0.25), 11.0 * 0.25 + 2.0 * (1.0 - 0.25)],
+          self.evaluate(var))
+    else:
+      assign = moving_averages.assign_moving_average(
+          var, val, decay, zero_debias=False)
+      self.evaluate(variables.global_variables_initializer())
+      self.assertAllClose([10.0, 11.0], self.evaluate(var))
+      assign.op.run()
+      self.assertAllClose(
+          [10.0 * 0.25 + 1.0 * (1.0 - 0.25), 11.0 * 0.25 + 2.0 * (1.0 - 0.25)],
+          self.evaluate(var))
+
+  @test_util.run_in_graph_and_eager_modes
+  def testAssignMovingAverage(self):
+    var = variables.Variable([0.0, 0.0])
+    val = constant_op.constant([1.0, 2.0], dtypes.float32)
+    decay = 0.25
+    if context.executing_eagerly():
+      self.assertAllClose([0.0, 0.0], self.evaluate(var))
+      assign = moving_averages.assign_moving_average(var, val, decay)
+      self.assertAllClose(
+          [1.0 * (1.0 - 0.25) / (1 - 0.25), 2.0 * (1.0 - 0.25) / (1 - 0.25)],
+          self.evaluate(var))
+    else:
+      assign = moving_averages.assign_moving_average(var, val, decay)
+      self.evaluate(variables.global_variables_initializer())
+      self.assertAllClose([0.0, 0.0], self.evaluate(var))
+      assign.op.run()
+      self.assertAllClose(
+          [1.0 * (1.0 - 0.25) / (1 - 0.25), 2.0 * (1.0 - 0.25) / (1 - 0.25)],
+          self.evaluate(var))
+
+  @test_util.deprecated_graph_mode_only
+  def testAssignMovingAverageNewNamingMultipleCalls(self):
+    with variable_scope.variable_scope("scope1") as vs1:
+      with variable_scope.variable_scope("scope2"):
+        var = variables.Variable(1.0, name="Var")
+        moving_averages.assign_moving_average(var, 0.0, 0.99)
+        moving_averages.assign_moving_average(var, 0.0, 0.99)
+    expected_names = ["scope1/scope2/Var:0",
+                      "scope1/scope2/scope1/scope2/Var/biased:0",
+                      "scope1/scope2/scope1/scope2/Var/local_step:0",
+                      "scope1/scope2/scope1/scope2/Var/biased_1:0",
+                      "scope1/scope2/scope1/scope2/Var/local_step_1:0"]
+    actual_names = [v.name for v in vs1.global_variables()]
+    self.assertSetEqual(set(expected_names), set(actual_names))
+
+  @test_util.deprecated_graph_mode_only
+  def testAssignMovingAverageNewNamingMultipleCallsWithReuse(self):
+    with variable_scope.variable_scope("scope1") as vs1:
+      var = variable_scope.get_variable("Var", shape=[])
+      moving_averages.assign_moving_average(var, 0.0, 0.99)
+      moving_averages.assign_moving_average(var, 0.0, 0.99)
+    with variable_scope.variable_scope(vs1, reuse=True):
+      var = variable_scope.get_variable("Var", shape=[])
+      moving_averages.assign_moving_average(var, 0.0, 0.99)
+      moving_averages.assign_moving_average(var, 0.0, 0.99)
+
+  @test_util.deprecated_graph_mode_only
+  def testWeightedMovingAverage(self):
+    with self.cached_session() as sess:
+      decay = 0.5
+      weight = array_ops.placeholder(dtypes.float32, [])
+      val = array_ops.placeholder(dtypes.float32, [])
+
+      wma = moving_averages.weighted_moving_average(val, decay, weight)
+      self.evaluate(variables.global_variables_initializer())
+
+      # Get the first weighted moving average.
+      val_1 = 3.0
+      weight_1 = 4.0
+      wma_array = sess.run(wma, feed_dict={val: val_1, weight: weight_1})
+      numerator_1 = val_1 * weight_1 * (1.0 - decay)
+      denominator_1 = weight_1 * (1.0 - decay)
+      self.assertAllClose(numerator_1 / denominator_1, wma_array)
+
+      # Get the second weighted moving average.
+      val_2 = 11.0
+      weight_2 = 22.0
+      wma_array = sess.run(wma, feed_dict={val: val_2, weight: weight_2})
+      numerator_2 = numerator_1 * decay + val_2 * weight_2 * (1.0 - decay)
+      denominator_2 = denominator_1 * decay + weight_2 * (1.0 - decay)
+      self.assertAllClose(numerator_2 / denominator_2, wma_array)
+
+  @test_util.deprecated_graph_mode_only
+  def testWeightedMovingAverageBfloat16(self):
+    with self.cached_session() as sess:
+      decay = 0.5
+      weight = array_ops.placeholder(dtypes.bfloat16, [])
+      val = array_ops.placeholder(dtypes.bfloat16, [])
+
+      wma = moving_averages.weighted_moving_average(val, decay, weight)
+      self.evaluate(variables.global_variables_initializer())
+
+      # Get the first weighted moving average.
+      val_1 = 3.0
+      weight_1 = 4.0
+      wma_array = sess.run(wma, feed_dict={val: val_1, weight: weight_1})
+      numerator_1 = val_1 * weight_1 * (1.0 - decay)
+      denominator_1 = weight_1 * (1.0 - decay)
+      self.assertAllClose(numerator_1 / denominator_1, wma_array)
+
+      # Get the second weighted moving average.
+      val_2 = 11.0
+      weight_2 = 22.0
+      wma_array = sess.run(wma, feed_dict={val: val_2, weight: weight_2})
+      numerator_2 = numerator_1 * decay + val_2 * weight_2 * (1.0 - decay)
+      denominator_2 = denominator_1 * decay + weight_2 * (1.0 - decay)
+      self.assertAllClose(
+          dtypes._np_bfloat16(numerator_2 / denominator_2), wma_array)
+
+
+def _Repeat(value, dim):
+  if dim == 1:
+    return value
+  return [value] * dim
+
+
+class ExponentialMovingAverageTest(test.TestCase):
+
+  def _CheckDecay(self, ema, actual_decay, dim, dynamic_decay_value=None):
+
+    def _Scale(dk, steps):
+      if ema._zero_debias:
+        return 1 - dk**steps
+      else:
+        return 1
+
+    tens = _Repeat(10.0, dim)
+    thirties = _Repeat(30.0, dim)
+    var0 = variables.Variable(tens, name="v0")
+    var1 = variables.Variable(thirties, name="v1")
+    self.evaluate(variables.global_variables_initializer())
+    # Note that tensor2 is not a Variable but just a plain Tensor resulting
+    # from the sum operation.
+    tensor2 = var0 + var1
+    if dynamic_decay_value is not None:
+      self.evaluate(ema._decay.assign(dynamic_decay_value))
+    update = ema.apply([var0, var1, tensor2])
+    avg0 = ema.average(var0)
+    avg1 = ema.average(var1)
+    avg2 = ema.average(tensor2)
+
+    self.assertItemsEqual([var0, var1], variables.moving_average_variables())
+
+    self.assertNotIn(avg0, variables.trainable_variables())
+    self.assertNotIn(avg1, variables.trainable_variables())
+    self.assertNotIn(avg2, variables.trainable_variables())
+    self.evaluate(variables.global_variables_initializer())
+    if dynamic_decay_value is not None:
+      self.evaluate(ema._decay.assign(dynamic_decay_value))
+
+    self.assertEqual("v0/ExponentialMovingAverage:0", avg0.name)
+    self.assertEqual("v1/ExponentialMovingAverage:0", avg1.name)
+    self.assertEqual("add/ExponentialMovingAverage:0", avg2.name)
+
+    # Check initial values.
+    self.assertAllClose(tens, self.evaluate(var0))
+    self.assertAllClose(thirties, self.evaluate(var1))
+    self.assertAllClose(_Repeat(10.0 + 30.0, dim), self.evaluate(tensor2))
+
+    # Check that averages are initialized correctly.
+    self.assertAllClose(tens, self.evaluate(avg0))
+    self.assertAllClose(thirties, self.evaluate(avg1))
+    # Note that averages of Tensor's initialize to zeros_like since no value
+    # of the Tensor is known because the Op has not been run (yet).
+    self.assertAllClose(_Repeat(0.0, dim), self.evaluate(avg2))
+
+    # Update the averages and check.
+    self.evaluate(update)
+    dk = actual_decay
+
+    expected = _Repeat(10.0 * dk + 10.0 * (1 - dk), dim)
+    self.assertAllClose(expected, self.evaluate(avg0))
+    expected = _Repeat(30.0 * dk + 30.0 * (1 - dk), dim)
+    self.assertAllClose(expected, self.evaluate(avg1))
+    expected = _Repeat(0.0 * dk + (10.0 + 30.0) * (1 - dk) / _Scale(dk, 1), dim)
+    self.assertAllClose(expected, self.evaluate(avg2))
+
+    # Again, update the averages and check.
+    self.evaluate(update)
+    expected = _Repeat((10.0 * dk + 10.0 * (1 - dk)) * dk + 10.0 * (1 - dk),
+                       dim)
+    self.assertAllClose(expected, self.evaluate(avg0))
+    expected = _Repeat((30.0 * dk + 30.0 * (1 - dk)) * dk + 30.0 * (1 - dk),
+                       dim)
+    self.assertAllClose(expected, self.evaluate(avg1))
+    expected = _Repeat(((0.0 * dk + (10.0 + 30.0) * (1 - dk)) * dk +
+                        (10.0 + 30.0) * (1 - dk)) / _Scale(dk, 2), dim)
+    self.assertAllClose(expected, self.evaluate(avg2))
+
+  @test_util.deprecated_graph_mode_only
+  def testAverageVariablesNoNumUpdates_Scalar(self):
+    ema = moving_averages.ExponentialMovingAverage(0.25)
+    self._CheckDecay(ema, actual_decay=0.25, dim=1)
+
+  @test_util.deprecated_graph_mode_only
+  def testAverageVariablesNoNumUpdates_Scalar_DynamicDecay(self):
+    decay_var = variables.Variable(0.75)
+    ema = moving_averages.ExponentialMovingAverage(decay_var)
+    self._CheckDecay(ema, actual_decay=0.25, dim=1, dynamic_decay_value=0.25)
+
+  @test_util.deprecated_graph_mode_only
+  def testAverageVariablesNoNumUpdates_Scalar_Debias(self):
+    ema = moving_averages.ExponentialMovingAverage(0.25, zero_debias=True)
+    self._CheckDecay(ema, actual_decay=0.25, dim=1)
+
+  @test_util.deprecated_graph_mode_only
+  def testAverageVariablesNoNumUpdates_Scalar_Debias_DynamicDecay(self):
+    decay_var = variables.Variable(0.75)
+    ema = moving_averages.ExponentialMovingAverage(decay_var, zero_debias=True)
+    self._CheckDecay(ema, actual_decay=0.25, dim=1, dynamic_decay_value=0.25)
+
+  @test_util.deprecated_graph_mode_only
+  def testAverageVariablesNoNumUpdates_Vector(self):
+    ema = moving_averages.ExponentialMovingAverage(0.25)
+    self._CheckDecay(ema, actual_decay=0.25, dim=5)
+
+  @test_util.deprecated_graph_mode_only
+  def testAverageVariablesNoNumUpdates_Vector_DynamicDecay(self):
+    decay_var = variables.Variable(0.75)
+    ema = moving_averages.ExponentialMovingAverage(decay_var)
+    self._CheckDecay(ema, actual_decay=0.25, dim=5, dynamic_decay_value=0.25)
+
+  @test_util.deprecated_graph_mode_only
+  def testAverageVariablesNoNumUpdates_Vector_Debias(self):
+    ema = moving_averages.ExponentialMovingAverage(0.25, zero_debias=True)
+    self._CheckDecay(ema, actual_decay=0.25, dim=5)
+
+  @test_util.deprecated_graph_mode_only
+  def testAverageVariablesNoNumUpdates_Vector_Debias_DynamicDecay(self):
+    decay_var = variables.Variable(0.75)
+    ema = moving_averages.ExponentialMovingAverage(decay_var, zero_debias=True)
+    self._CheckDecay(ema, actual_decay=0.25, dim=5, dynamic_decay_value=0.25)
+
+  @test_util.deprecated_graph_mode_only
+  def testAverageVariablesNumUpdates_Scalar(self):
+    # With num_updates 1, the decay applied is 0.181818.
+    ema = moving_averages.ExponentialMovingAverage(0.25, num_updates=1)
+    self._CheckDecay(ema, actual_decay=0.181818, dim=1)
+
+  @test_util.deprecated_graph_mode_only
+  def testAverageVariablesNumUpdates_Scalar_DynamicDecay(self):
+    decay_var = variables.Variable(0.75)
+    # With num_updates 1, the decay applied is 0.181818.
+    ema = moving_averages.ExponentialMovingAverage(decay_var, num_updates=1)
+    self._CheckDecay(
+        ema, actual_decay=0.181818, dim=1, dynamic_decay_value=0.25)
+
+  @test_util.deprecated_graph_mode_only
+  def testAverageVariablesNumUpdates_Scalar_Debias(self):
+    # With num_updates 1, the decay applied is 0.181818.
+    ema = moving_averages.ExponentialMovingAverage(
+        0.25, num_updates=1, zero_debias=True)
+    self._CheckDecay(ema, actual_decay=0.181818, dim=1)
+
+  @test_util.deprecated_graph_mode_only
+  def testAverageVariablesNumUpdates_Scalar_Debias_DynamicDecay(self):
+    decay_var = variables.Variable(0.75)
+    # With num_updates 1, the decay applied is 0.181818.
+    ema = moving_averages.ExponentialMovingAverage(
+        decay_var, num_updates=1, zero_debias=True)
+    self._CheckDecay(
+        ema, actual_decay=0.181818, dim=1, dynamic_decay_value=0.25)
+
+  @test_util.deprecated_graph_mode_only
+  def testAverageVariablesNumUpdates_Vector(self):
+    # With num_updates 1, the decay applied is 0.181818.
+    ema = moving_averages.ExponentialMovingAverage(0.25, num_updates=1)
+    self._CheckDecay(ema, actual_decay=0.181818, dim=5)
+
+  @test_util.deprecated_graph_mode_only
+  def testAverageVariablesNumUpdates_Vector_DynamicDecay(self):
+    decay_var = variables.Variable(0.75)
+    # With num_updates 1, the decay applied is 0.181818.
+    ema = moving_averages.ExponentialMovingAverage(decay_var, num_updates=1)
+    self._CheckDecay(
+        ema, actual_decay=0.181818, dim=5, dynamic_decay_value=0.25)
+
+  @test_util.deprecated_graph_mode_only
+  def testAverageVariablesNumUpdates_Vector_Debias(self):
+    # With num_updates 1, the decay applied is 0.181818.
+    ema = moving_averages.ExponentialMovingAverage(
+        0.25, num_updates=1, zero_debias=True)
+    self._CheckDecay(ema, actual_decay=0.181818, dim=5)
+
+  @test_util.deprecated_graph_mode_only
+  def testAverageVariablesNumUpdates_Vector_Debias_DynamicDecay(self):
+    decay_var = variables.Variable(0.75)
+    # With num_updates 1, the decay applied is 0.181818.
+    ema = moving_averages.ExponentialMovingAverage(
+        decay_var, num_updates=1, zero_debias=True)
+    self._CheckDecay(
+        ema, actual_decay=0.181818, dim=5, dynamic_decay_value=0.25)
+
+  @test_util.deprecated_graph_mode_only
+  def testAverageVariablesWithControlDeps(self):
+    v0 = variables.Variable(0, name="v0")
+    add_to_v0 = v0.assign_add(1)
+    v1 = variables.Variable([10.0], name="v1")
+    assign_to_v1 = v1.assign([20.0])
+    ema = moving_averages.ExponentialMovingAverage(0.25)
+    with ops.control_dependencies([add_to_v0]):
+      ema_op = ema.apply([v1])
+    # the moving average of v1 should not have any control inputs
+    v1_avg = ema.average(v1)
+    self.assertEqual([], v1_avg.initializer.control_inputs)
+    self.assertEqual([], v1_avg.value().op.control_inputs)
+    self.assertEqual([], v1_avg.value().op.control_inputs)
+    # We should be able to initialize v1_avg before v0.
+    self.evaluate(v1_avg.initializer)
+    self.evaluate(v0.initializer)
+    self.assertEqual([10.0], self.evaluate(v1_avg))
+    # running ema_op should add to v0 (in addition to updating v1_avg)
+    self.evaluate(assign_to_v1)
+    self.evaluate(ema_op)
+    self.assertEqual(1, self.evaluate(v0))
+    self.assertEqual([17.5], self.evaluate(v1_avg))
+
+  def testBasicEager(self):
+    v0 = variables.Variable(1.0, name="v0")
+    v1 = variables.Variable(2.0, name="v1")
+
+    ema = moving_averages.ExponentialMovingAverage(0.25, name="foo")
+    op = ema.apply([v0, v1])
+    if not context.executing_eagerly():
+      self.evaluate(variables.global_variables_initializer())
+      self.evaluate(op)
+
+    self.evaluate(v0.assign(2.0))
+    self.evaluate(v1.assign(4.0))
+
+    self.evaluate(ema.apply([v0, v1]))
+
+    self.assertEqual("foo", ema.name)
+    self.assertEqual("v0/foo", ema.average_name(v0))
+    self.assertEqual("v1/foo", ema.average_name(v1))
+
+    self.assertAllEqual(self.evaluate(ema.average(v0)), 1.75)
+    self.assertAllEqual(self.evaluate(ema.average(v1)), 3.5)
+
+  def averageVariablesNamesHelper(self, zero_debias):
+    v0 = variables.Variable(10.0, name="v0")
+    v1 = variables.Variable(30.0, name="v1")
+    # Add a non-trainable variable.
+    v2 = variables.Variable(20.0, name="v2", trainable=False)
+    tensor2 = v0 + v1
+    ema = moving_averages.ExponentialMovingAverage(
+        0.25, zero_debias=zero_debias, name="foo")
+    self.assertEqual("foo", ema.name)
+    self.assertEqual("v0/foo", ema.average_name(v0))
+    self.assertEqual("v1/foo", ema.average_name(v1))
+    self.assertEqual("add/foo", ema.average_name(tensor2))
+    ema.apply([v0, v1, tensor2])
+    vars_to_restore = ema.variables_to_restore()
+    # vars_to_restore should contain the following:
+    # {v0/foo : v0,
+    #  v1/foo : v1,
+    #  add/foo : add/foo,
+    #  v2 : v2}
+    expected_names = [
+        ema.average_name(v0),
+        ema.average_name(v1),
+        ema.average_name(tensor2), v2.op.name
+    ]
+    if zero_debias:
+      # vars_to_restore should also contain the following:
+      #  {add/foo/biased: add/foo/biased,
+      #  add/foo/local_step: add/foo/local_step}
+      expected_names += [
+          ema.average_name(tensor2) + "/biased",
+          ema.average_name(tensor2) + "/local_step"
+      ]
+    self.assertEqual(sorted(expected_names), sorted(vars_to_restore.keys()))
+    self.assertEqual(ema.average(v0).op.name, ema.average_name(v0))
+    self.assertEqual(ema.average(v1).op.name, ema.average_name(v1))
+    self.assertEqual(ema.average(tensor2).op.name, ema.average_name(tensor2))
+
+  @test_util.deprecated_graph_mode_only
+  def testAverageVariablesNames(self):
+    self.averageVariablesNamesHelper(zero_debias=True)
+
+  @test_util.deprecated_graph_mode_only
+  def testAverageVariablesNamesNoDebias(self):
+    self.averageVariablesNamesHelper(zero_debias=False)
+
+  @test_util.deprecated_graph_mode_only
+  def averageVariablesNamesRespectScopeHelper(self, zero_debias):
+    # See discussion on #2740.
+    with variable_scope.variable_scope("scope1"):
+      v0 = variables.Variable(10.0, name="v0")
+      v1 = variables.Variable(30.0, name="v1")
+      # Add a non-trainable variable.
+      v2 = variables.Variable(20.0, name="v2", trainable=False)
+      tensor2 = v0 + v1
+    with variable_scope.variable_scope("scope2"):
+      ema = moving_averages.ExponentialMovingAverage(
+          0.25, zero_debias=zero_debias, name="foo")
+      self.assertEqual("scope2/scope1/v0/foo", ema.average_name(v0))
+      self.assertEqual("scope2/scope1/v1/foo", ema.average_name(v1))
+      self.assertEqual("scope2/scope1/add/foo", ema.average_name(tensor2))
+      ema.apply([v0, v1, tensor2])
+      vars_to_restore = ema.variables_to_restore()
+      # `vars_to_restore` should contain the following:
+      # {scope2/scope1/v0/foo : v0,
+      #  scope2/scope1/v1/foo : v1,
+      #  scope2/scope1/add/foo : add/foo,
+      #  scope1/v2 : v2}
+      expected_names = [
+          ema.average_name(v0),
+          ema.average_name(v1),
+          ema.average_name(tensor2), v2.op.name
+      ]
+      if zero_debias:
+        # `vars_to_restore` should also contain the following:
+        # {scope2/scope2/scope1/add/foo/biased: add/foo/biased,
+        #  scope2/scope2/scope1/add/foo/local_step: add/foo/local_step}
+        sc = "scope2/"
+        expected_names += [
+            sc + ema.average_name(tensor2) + "/biased",
+            sc + ema.average_name(tensor2) + "/local_step"
+        ]
+
+      self.assertEqual(sorted(expected_names), sorted(vars_to_restore.keys()))
+      self.assertEqual(ema.average(v0).op.name, ema.average_name(v0))
+      self.assertEqual(ema.average(v1).op.name, ema.average_name(v1))
+      self.assertEqual(ema.average(tensor2).op.name, ema.average_name(tensor2))
+
+  @test_util.deprecated_graph_mode_only
+  def testAverageVariablesNamesRespectScope(self):
+    self.averageVariablesNamesRespectScopeHelper(zero_debias=True)
+
+  @test_util.deprecated_graph_mode_only
+  def testAverageVariablesNamesRespectScopeNoDebias(self):
+    self.averageVariablesNamesRespectScopeHelper(zero_debias=False)
+
+  @test_util.deprecated_graph_mode_only
+  def testSubsetAverageVariablesNames(self):
+    v0 = variables.Variable(10.0, name="v0")
+    v1 = variables.Variable(30.0, name="v1")
+    # Add a non-trainable variable.
+    v2 = variables.Variable(20.0, name="v2", trainable=False)
+    tensor2 = v0 + v1
+    ema = moving_averages.ExponentialMovingAverage(0.25, name="foo_avg")
+    self.assertEqual("v0/foo_avg", ema.average_name(v0))
+    self.assertEqual("v1/foo_avg", ema.average_name(v1))
+    self.assertEqual("add/foo_avg", ema.average_name(tensor2))
+    vars_to_restore = ema.variables_to_restore([v0, tensor2])
+    # vars_to_restore should contain the following:
+    # {v0/foo_avg : v0,
+    #  add/foo_avg : add
+    #  v1 : v1,
+    #  v2 : v2}
+    self.assertEqual(
+        sorted(vars_to_restore.keys()),
+        sorted([
+            ema.average_name(v0),
+            ema.average_name(tensor2), v1.op.name, v2.op.name
+        ]))
+    ema.apply([v0, v1, tensor2])
+    self.assertEqual(ema.average(v0).op.name, ema.average_name(v0))
+    self.assertEqual(ema.average(v1).op.name, ema.average_name(v1))
+    self.assertEqual(ema.average(tensor2).op.name, ema.average_name(tensor2))
+
+  def testSubsetAverageVariablesNamesEager(self):
+    v0 = variables.Variable(10.0, name="v0")
+    v1 = variables.Variable(30.0, name="v1")
+    # Add a non-trainable variable.
+    v2 = variables.Variable(20.0, name="v2", trainable=False)
+    ema = moving_averages.ExponentialMovingAverage(0.25, name="foo_avg")
+    self.assertEqual("v0/foo_avg", ema.average_name(v0))
+    self.assertEqual("v1/foo_avg", ema.average_name(v1))
+    vars_to_restore = ema.variables_to_restore([v0, v1, v2])
+    self.assertAllEqual(
+        sorted(vars_to_restore.keys()),
+        sorted([
+            ema.average_name(v0), ema.average_name(v1), ema.average_name(v2)
+        ]))
+    ema.apply([v0, v1])
+    self.assertEqual(ema.average(v0).name[:-len(":0")], ema.average_name(v0))
+    self.assertEqual(ema.average(v1).name[:-len(":0")], ema.average_name(v1))
+
+  @test_util.deprecated_graph_mode_only
+  def testAverageVariablesDeviceAssignment(self):
+    with ops.device("/job:dev_v0"):
+      v0 = variables.Variable(10.0, name="v0")
+    with ops.device("/job:dev_v1"):
+      v1 = gen_state_ops.variable(
+          shape=[1],
+          dtype=dtypes.float32,
+          name="v1",
+          container="",
+          shared_name="")
+      v1.set_shape([1])
+    tensor2 = v0 + v1
+    ema = moving_averages.ExponentialMovingAverage(0.25, name="foo_avg")
+    with ops.device("/job:default"):
+      ema.apply([v0, v1, tensor2])
+    self.assertDeviceEqual("/job:dev_v0", ema.average(v0).device)
+    self.assertDeviceEqual("/job:dev_v1", ema.average(v1).device)
+    # However, the colocation property is maintained.
+    self.assertEqual([b"loc:@v1"], ema.average(v1).op.colocation_groups())
+    self.assertDeviceEqual("/job:default", ema.average(tensor2).device)
+
+  def _ExportAndImportGraph(self, graph):
+    """Export and import graph into a new graph."""
+    meta_graph = saver_lib.export_meta_graph(
+        graph=graph, collection_list=graph.get_all_collection_keys())
+    graph_copy = ops.Graph()
+    with graph_copy.as_default():
+      _ = saver_lib.import_meta_graph(meta_graph)
+    return graph_copy
+
+  @test_util.deprecated_graph_mode_only
+  def testImportedGraphVariablesToRestore(self):
+    g = ops.Graph()
+    with g.as_default():
+      variables.Variable(10.0, name="v")
+    # Export and import the graph into a new graph.
+    g_copy = self._ExportAndImportGraph(g)
+    with g_copy.as_default():
+      ema = moving_averages.ExponentialMovingAverage(0.25, name="foo_avg")
+      vars_to_restore = ema.variables_to_restore()
+      # There should only be one variable in vars_to_restore. This is important
+      # to check because when importing from a GraphDef, TF makes duplicate
+      # python Variable objects referring to the same underlying variable. We
+      # need to be sure that two variables referring to the same variable don't
+      # both get added to vars_to_restore.
+      self.assertEqual(len(vars_to_restore), 1)
+      self.assertIn("v/foo_avg", vars_to_restore)
+
+  @test_util.deprecated_graph_mode_only
+  def testCopyXlaSharding(self):
+    ema = moving_averages.ExponentialMovingAverage(0.25, name="foo_avg")
+    v = variables.Variable(_Repeat(10.0, 2), name="v")
+    self.assertIsNone(xla_sharding.get_tensor_sharding(v))
+    v = xla_sharding.mesh_split(v, np.array([0, 1]), [0], use_sharding_op=False)
+    self.assertIsNotNone(xla_sharding.get_tensor_sharding(v))
+    self.evaluate(variables.global_variables_initializer())
+    ema.apply([v])
+    avg = ema.average(v)
+    self.assertEqual(
+        xla_sharding.get_tensor_sharding(v),
+        xla_sharding.get_tensor_sharding(avg))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/optimizer_test.py

@@ -0,0 +1,296 @@
+# 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 optimizer."""
+
+from tensorflow.python.distribute import cross_device_ops
+from tensorflow.python.distribute import distribute_utils
+from tensorflow.python.distribute import mirrored_strategy
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import clip_ops
+from tensorflow.python.ops import gradients_util
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import state_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import adam
+from tensorflow.python.training import gradient_descent
+
+
+class OptimizerTest(test.TestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def testBasic(self):
+    for i, dtype in enumerate([dtypes.half, dtypes.float32, dtypes.float64]):
+      # Note that we name the variables uniquely here since the variables don't
+      # seem to be getting deleted at the end of the loop.
+      var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype,
+                                                    name='a_%d' % i)
+      var1 = resource_variable_ops.ResourceVariable([3.0, 4.0], dtype=dtype,
+                                                    name='b_%d' % i)
+      def loss():
+        return 5 * var0 + 3 * var1  # pylint: disable=cell-var-from-loop
+      # Note that for eager execution, minimize expects a function instead of a
+      # Tensor.
+      global_step = resource_variable_ops.ResourceVariable(
+          array_ops.zeros([], dtypes.int64), name='global_step_%d' % i)
+      sgd_op = gradient_descent.GradientDescentOptimizer(3.0)
+
+      self.evaluate(variables.global_variables_initializer())
+      # Fetch params to validate initial values
+      self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+      self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+      # Run 1 step of sgd through optimizer
+      opt_op = sgd_op.minimize(loss, global_step, [var0, var1])
+      self.evaluate(opt_op)
+      # Validate updated params
+      self.assertAllClose([-14., -13.], self.evaluate(var0))
+      self.assertAllClose([-6., -5.], self.evaluate(var1))
+
+  @test_util.run_deprecated_v1
+  def testAggregationMethod(self):
+    for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+      with self.cached_session():
+        var0 = variables.Variable([1.0, 2.0], dtype=dtype)
+        var1 = variables.Variable([3.0, 4.0], dtype=dtype)
+        cost = 5 * var0 + 3 * var1
+        global_step = variables.Variable(
+            array_ops.zeros([], dtypes.int64), name='global_step')
+        sgd_op = gradient_descent.GradientDescentOptimizer(3.0)
+        opt_op = sgd_op.minimize(
+            cost,
+            global_step, [var0, var1],
+            aggregation_method=gradients_util.AggregationMethod.
+            EXPERIMENTAL_ACCUMULATE_N)
+
+        self.evaluate(variables.global_variables_initializer())
+        # Fetch params to validate initial values
+        self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+        self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+        # Run 1 step of sgd through optimizer
+        opt_op.run()
+        # Validate updated params
+        self.assertAllClose([-14., -13.], self.evaluate(var0))
+        self.assertAllClose([-6., -5.], self.evaluate(var1))
+
+  @test_util.run_deprecated_v1
+  def testPrecomputedGradient(self):
+    for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+      with self.cached_session():
+        var0 = variables.Variable([1.0, 2.0], dtype=dtype)
+        var1 = variables.Variable([3.0, 4.0], dtype=dtype)
+        cost = 5 * var0 + 3 * var1
+        grad_loss = constant_op.constant([42, -42], dtype=dtype)
+        global_step = variables.Variable(
+            array_ops.zeros([], dtypes.int64), name='global_step')
+        sgd_op = gradient_descent.GradientDescentOptimizer(3.0)
+        opt_op = sgd_op.minimize(
+            cost, global_step, [var0, var1], grad_loss=grad_loss)
+
+        self.evaluate(variables.global_variables_initializer())
+        # Fetch params to validate initial values
+        self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+        self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+        # Run 1 step of sgd through optimizer
+        opt_op.run()
+        # Validate updated params
+        self.assertAllClose([1.0 - 3 * 5 * 42.0, 2.0 - 3 * 5 * (-42.0)],
+                            self.evaluate(var0))
+        self.assertAllClose([3.0 - 3 * 3 * 42.0, 4.0 - 3 * 3 * (-42.0)],
+                            self.evaluate(var1))
+
+  @test_util.run_in_graph_and_eager_modes
+  def testNoVariables(self):
+    for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+      # pylint: disable=cell-var-from-loop
+      def loss():
+        var0 = resource_variable_ops.ResourceVariable(
+            [1.0, 2.0], dtype=dtype, trainable=False, name='a')
+        var1 = resource_variable_ops.ResourceVariable(
+            [3.0, 4.0], dtype=dtype, trainable=False, name='b')
+        return 5 * var0 + var1
+      # pylint: enable=cell-var-from-loop
+      sgd_op = gradient_descent.GradientDescentOptimizer(3.0)
+      with self.assertRaisesRegex(ValueError, 'No.*variables'):
+        sgd_op.minimize(loss)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testNoGradients(self):
+    for i, dtype in enumerate([dtypes.half, dtypes.float32, dtypes.float64]):
+      # Note that we name the variables uniquely here since the variables don't
+      # seem to be getting deleted at the end of the loop.
+      var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype,
+                                                    name='a%d' % i)
+      var1 = resource_variable_ops.ResourceVariable([3.0, 4.0], dtype=dtype,
+                                                    name='b%d' % i)
+      # pylint: disable=cell-var-from-loop
+      def loss():
+        return 5 * var0
+      # pylint: enable=cell-var-from-loop
+      sgd_op = gradient_descent.GradientDescentOptimizer(3.0)
+      with self.assertRaisesRegex(ValueError, 'No gradients'):
+        # var1 has no gradient
+        sgd_op.minimize(loss, var_list=[var1])
+
+  @test_util.run_in_graph_and_eager_modes
+  def testNoGradientsForAnyVariables_Minimize(self):
+    for i, dtype in enumerate([dtypes.half, dtypes.float32, dtypes.float64]):
+      # Note that we name the variables uniquely here since the variables don't
+      # seem to be getting deleted at the end of the loop.
+      var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype,
+                                                    name='a_%d' % i)
+      var1 = resource_variable_ops.ResourceVariable([3.0, 4.0], dtype=dtype,
+                                                    name='b_%d' % i)
+      def loss():
+        return constant_op.constant(5.0)
+      sgd_op = gradient_descent.GradientDescentOptimizer(3.0)
+      with self.assertRaisesRegex(ValueError,
+                                  'No gradients provided for any variable'):
+        sgd_op.minimize(loss, var_list=[var0, var1])
+
+  @test_util.run_in_graph_and_eager_modes
+  def testNoGradientsForAnyVariables_ApplyGradients(self):
+    for i, dtype in enumerate([dtypes.half, dtypes.float32, dtypes.float64]):
+      # Note that we name the variables uniquely here since the variables don't
+      # seem to be getting deleted at the end of the loop.
+      var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype,
+                                                    name='a_%d' % i)
+      var1 = resource_variable_ops.ResourceVariable([3.0, 4.0], dtype=dtype,
+                                                    name='b_%d' % i)
+      sgd_op = gradient_descent.GradientDescentOptimizer(3.0)
+      with self.assertRaisesRegex(ValueError,
+                                  'No gradients provided for any variable'):
+        sgd_op.apply_gradients([(None, var0), (None, var1)])
+
+  @test_util.run_in_graph_and_eager_modes
+  def testGradientsAsVariables(self):
+    for i, dtype in enumerate([dtypes.half, dtypes.float32, dtypes.float64]):
+      # Note that we name the variables uniquely here since the variables don't
+      # seem to be getting deleted at the end of the loop.
+      var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype,
+                                                    name='a%d' % i)
+      var1 = resource_variable_ops.ResourceVariable([3.0, 4.0], dtype=dtype,
+                                                    name='b%d' % i)
+      def loss():
+        return 5 * var0 + 3 * var1  # pylint: disable=cell-var-from-loop
+      sgd_op = gradient_descent.GradientDescentOptimizer(3.0)
+      grads_and_vars = sgd_op.compute_gradients(loss, [var0, var1])
+      # Convert gradients to tf.Variables
+      converted_grads = [
+          resource_variable_ops.ResourceVariable(array_ops.zeros([2], dtype),
+                                                 name='c_%d_%d' % (i, j))
+          for j, gv in enumerate(grads_and_vars)
+      ]
+      convert_ops = [
+          state_ops.assign(converted_grads[j], gv[0])
+          for j, gv in enumerate(grads_and_vars)
+      ]
+
+      self.evaluate(variables.global_variables_initializer())
+      # Run convert_ops to achieve the gradients converting
+      self.evaluate(convert_ops)
+      # Fetch params to validate initial values
+      self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+      self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+
+      # Run 1 step of sgd through optimizer
+      converted_grads_and_vars = list(zip(converted_grads, [var0, var1]))
+      opt_op = sgd_op.apply_gradients(converted_grads_and_vars)
+      self.evaluate(opt_op)
+
+      # Validate updated params
+      self.assertAllClose([-14., -13.], self.evaluate(var0))
+      self.assertAllClose([-6., -5.], self.evaluate(var1))
+
+  @test_util.run_in_graph_and_eager_modes
+  def testComputeGradientsWithTensors(self):
+    x = ops.convert_to_tensor(1.0)
+    def f():
+      return x * x
+    sgd_op = gradient_descent.GradientDescentOptimizer(3.0)
+    grads_and_vars = sgd_op.compute_gradients(f, [x])
+    self.assertEqual(1, len(grads_and_vars))
+    grad, x_as_var = grads_and_vars[0]
+    self.assertIs(x, x_as_var)
+    self.assertEqual(2.0, self.evaluate(grad))
+
+    with self.assertRaises(NotImplementedError):
+      sgd_op.apply_gradients(grads_and_vars)
+
+  @test_util.run_deprecated_v1
+  def testTrainOp(self):
+    with self.cached_session():
+      var0 = variables.Variable([1.0, 2.0])
+      var1 = variables.Variable([3.0, 4.0])
+      cost = 5 * var0 + 3 * var1
+      global_step = variables.Variable(
+          array_ops.zeros([], dtypes.int64), name='global_step')
+      sgd_op = gradient_descent.GradientDescentOptimizer(3.0)
+      opt_op = sgd_op.minimize(cost, global_step, [var0, var1])
+      self.assertTrue(opt_op in ops.get_collection(ops.GraphKeys.TRAIN_OP))
+
+  @test_util.run_deprecated_v1
+  def testConstraint(self):
+    constraint_01 = lambda x: clip_ops.clip_by_value(x, -0.1, 0.)
+    constraint_0 = lambda x: clip_ops.clip_by_value(x, 0., 1.)
+    with self.cached_session():
+      var0 = variables.Variable([1.0, 2.0],
+                                constraint=constraint_01)
+      var1 = variables.Variable([3.0, 4.0],
+                                constraint=constraint_0)
+      cost = 5 * var0 + 3 * var1
+      global_step = variables.Variable(
+          array_ops.zeros([], dtypes.int64), name='global_step')
+      sgd_op = gradient_descent.GradientDescentOptimizer(3.0)
+      opt_op = sgd_op.minimize(cost, global_step, [var0, var1])
+
+      self.evaluate(variables.global_variables_initializer())
+      # Fetch params to validate initial values
+      self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+      self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+      # Run 1 step of sgd through optimizer
+      opt_op.run()
+      # Validate updated params
+      self.assertAllClose([-0.1, -0.1], self.evaluate(var0))
+      self.assertAllClose([0., 0.], self.evaluate(var1))
+
+  @test_util.run_deprecated_v1
+  def testGetSlotUnderDistributedStrategy(self):
+    # Only run this test in graph mode so we don't need actual GPU.
+    ds = mirrored_strategy.MirroredStrategy(
+        ['CPU:0', 'GPU:0'],
+        cross_device_ops=cross_device_ops.HierarchicalCopyAllReduce())
+    # We need an optimizer that creates slots.
+    optimizer = adam.AdamOptimizer()
+
+    def f():
+      v = variables.Variable([1.0])
+      self.assertTrue(distribute_utils.is_distributed_variable(v))
+      # Slot variables are created in the first call to apply_gradients.
+      optimizer.apply_gradients([(ops.convert_to_tensor([1.0]), v)])
+      self.assertTrue(optimizer.get_slot_names())
+      for name in optimizer.get_slot_names():
+        slot = optimizer.get_slot(v, name)
+        self.assertIsNotNone(slot)
+        self.assertTrue(distribute_utils.is_distributed_variable(slot))
+
+    ds.run(f)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/training/proximal_adagrad.py

@@ -0,0 +1,123 @@
+# 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.
+# ==============================================================================
+
+"""ProximalAdagrad for TensorFlow."""
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import gen_training_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.training import optimizer
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(v1=["train.ProximalAdagradOptimizer"])
+class ProximalAdagradOptimizer(optimizer.Optimizer):
+  # pylint: disable=line-too-long
+  """Optimizer that implements the Proximal Adagrad algorithm.
+
+  References:
+    Adaptive Subgradient Methods for Online Learning and Stochastic Optimization:
+      [Duchi et al., 2011](http://jmlr.org/papers/v12/duchi11a.html)
+      ([pdf](http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf))
+    Efficient Learning using Forward-Backward Splitting:
+      [Duchi et al., 2009](http://papers.nips.cc/paper/3793-efficient-learning-using-forward-backward-splitting)
+      ([pdf](http://papers.nips.cc/paper/3793-efficient-learning-using-forward-backward-splitting.pdf))
+  """
+
+  def __init__(self, learning_rate, initial_accumulator_value=0.1,
+               l1_regularization_strength=0.0, l2_regularization_strength=0.0,
+               use_locking=False, name="ProximalAdagrad"):
+    """Construct a new ProximalAdagrad optimizer.
+
+    Args:
+      learning_rate: A `Tensor` or a floating point value.  The learning rate.
+      initial_accumulator_value: A floating point value.
+        Starting value for the accumulators, must be positive.
+      l1_regularization_strength: A float value, must be greater than or
+        equal to zero.
+      l2_regularization_strength: A float value, must be greater than or
+        equal to zero.
+      use_locking: If `True` use locks for update operations.
+      name: Optional name prefix for the operations created when applying
+        gradients.  Defaults to "Adagrad".
+
+    Raises:
+      ValueError: If the `initial_accumulator_value` is invalid.
+    """
+    if initial_accumulator_value <= 0.0:
+      raise ValueError("initial_accumulator_value must be positive: %s" %
+                       initial_accumulator_value)
+    super(ProximalAdagradOptimizer, self).__init__(use_locking, name)
+    self._learning_rate = learning_rate
+    self._initial_accumulator_value = initial_accumulator_value
+    self._l1_regularization_strength = l1_regularization_strength
+    self._l2_regularization_strength = l2_regularization_strength
+    # Created in Initialize.
+    self._l1_regularization_strength_tensor = None
+    self._l2_regularization_strength_tensor = None
+    self._learning_rate_tensor = None
+
+  def _create_slots(self, var_list):
+    for v in var_list:
+      with ops.colocate_with(v):
+        val = constant_op.constant(self._initial_accumulator_value,
+                                   shape=v.get_shape(),
+                                   dtype=v.dtype.base_dtype)
+      self._get_or_make_slot(v, val, "accumulator", self._name)
+
+  def _prepare(self):
+    self._learning_rate_tensor = ops.convert_to_tensor(self._learning_rate,
+                                                       name="learning_rate")
+    self._l1_regularization_strength_tensor = ops.convert_to_tensor(
+        self._l1_regularization_strength,
+        name="l1_regularization_strength")
+    self._l2_regularization_strength_tensor = ops.convert_to_tensor(
+        self._l2_regularization_strength,
+        name="l2_regularization_strength")
+
+  def _apply_dense(self, grad, var):
+    acc = self.get_slot(var, "accumulator")
+    return gen_training_ops.apply_proximal_adagrad(
+        var, acc, self._learning_rate_tensor,
+        self._l1_regularization_strength_tensor,
+        self._l2_regularization_strength_tensor,
+        grad, use_locking=self._use_locking)
+
+  def _resource_apply_dense(self, grad, var):
+    acc = self.get_slot(var, "accumulator")
+    return gen_training_ops.resource_apply_proximal_adagrad(
+        var.handle, acc.handle, self._learning_rate_tensor,
+        self._l1_regularization_strength_tensor,
+        self._l2_regularization_strength_tensor,
+        grad, use_locking=self._use_locking)
+
+  def _apply_sparse(self, grad, var):
+    acc = self.get_slot(var, "accumulator")
+    return gen_training_ops.sparse_apply_proximal_adagrad(
+        var, acc, self._learning_rate_tensor,
+        self._l1_regularization_strength_tensor,
+        self._l2_regularization_strength_tensor,
+        grad.values, grad.indices,
+        use_locking=self._use_locking)
+
+  def _resource_apply_sparse(self, grad, var, indices):
+    acc = self.get_slot(var, "accumulator")
+    return gen_training_ops.resource_sparse_apply_proximal_adagrad(
+        var.handle, acc.handle,
+        math_ops.cast(self._learning_rate_tensor, grad.dtype),
+        math_ops.cast(self._l1_regularization_strength_tensor, grad.dtype),
+        math_ops.cast(self._l2_regularization_strength_tensor, grad.dtype),
+        grad, indices,
+        use_locking=self._use_locking)

tensorflow/python/training/proximal_adagrad_test.py

@@ -0,0 +1,251 @@
+# 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 tests for Proximal Adagrad operations."""
+
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import embedding_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import adagrad
+from tensorflow.python.training import proximal_adagrad
+
+
+class ProximalAdagradOptimizerTest(test.TestCase):
+
+  def doTestProximalAdagradwithoutRegularization(self, use_resource=False):
+    # ProximalAdagradOptimizer is supported only in V1.
+    with ops.Graph().as_default(), self.cached_session():
+      var0 = variables.Variable([0.0, 0.0])
+      var1 = variables.Variable([0.0, 0.0])
+      grads0 = constant_op.constant([0.1, 0.2])
+      grads1 = constant_op.constant([0.01, 0.02])
+      opt = proximal_adagrad.ProximalAdagradOptimizer(
+          3.0,
+          initial_accumulator_value=0.1,
+          l1_regularization_strength=0.0,
+          l2_regularization_strength=0.0)
+      update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+      self.evaluate(variables.global_variables_initializer())
+
+      v0_val, v1_val = self.evaluate([var0, var1])
+      self.assertAllClose([0.0, 0.0], v0_val)
+      self.assertAllClose([0.0, 0.0], v1_val)
+
+      # Run 3 steps Proximal Adagrad.
+      for _ in range(3):
+        update.run()
+
+      v0_val, v1_val = self.evaluate([var0, var1])
+      self.assertAllClose(np.array([-2.60260963, -4.29698515]), v0_val)
+      self.assertAllClose(np.array([-0.28432083, -0.56694895]), v1_val)
+      opt_vars = opt.variables()
+      self.assertStartsWith(opt_vars[0].name, var0._shared_name)
+      self.assertStartsWith(opt_vars[1].name, var1._shared_name)
+      self.assertEqual(2, len(opt_vars))
+
+  def testProximalAdagradwithoutRegularization(self):
+    self.doTestProximalAdagradwithoutRegularization(use_resource=False)
+
+  def testResourceProximalAdagradwithoutRegularization(self):
+    self.doTestProximalAdagradwithoutRegularization(use_resource=True)
+
+  def testProximalAdagradwithoutRegularization2(self):
+    # ProximalAdagradOptimizer is supported only in V1.
+    with ops.Graph().as_default(), self.cached_session():
+      var0 = variables.Variable([1.0, 2.0])
+      var1 = variables.Variable([4.0, 3.0])
+      grads0 = constant_op.constant([0.1, 0.2])
+      grads1 = constant_op.constant([0.01, 0.02])
+
+      opt = proximal_adagrad.ProximalAdagradOptimizer(
+          3.0,
+          initial_accumulator_value=0.1,
+          l1_regularization_strength=0.0,
+          l2_regularization_strength=0.0)
+      update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+      self.evaluate(variables.global_variables_initializer())
+
+      v0_val, v1_val = self.evaluate([var0, var1])
+      self.assertAllClose([1.0, 2.0], v0_val)
+      self.assertAllClose([4.0, 3.0], v1_val)
+
+      # Run 3 steps Proximal Adagrad.
+      for _ in range(3):
+        update.run()
+      v0_val, v1_val = self.evaluate([var0, var1])
+      self.assertAllClose(np.array([-1.60261, -2.296985]), v0_val)
+      self.assertAllClose(np.array([3.715679, 2.433051]), v1_val)
+
+  def testMinimizeSparseResourceVariable(self):
+    for dtype in [dtypes.float32, dtypes.float64]:
+      # ProximalAdagradOptimizer is supported only in V1.
+      with ops.Graph().as_default(), self.cached_session():
+        var0 = resource_variable_ops.ResourceVariable([[1.0, 2.0]], dtype=dtype)
+        x = constant_op.constant([[4.0], [5.0]], dtype=dtype)
+        pred = math_ops.matmul(embedding_ops.embedding_lookup([var0], [0]), x)
+        loss = pred * pred
+        sgd_op = proximal_adagrad.ProximalAdagradOptimizer(1.0).minimize(loss)
+        self.evaluate(variables.global_variables_initializer())
+        # Fetch params to validate initial values
+        self.assertAllCloseAccordingToType([[1.0, 2.0]], self.evaluate(var0))
+        # Run 1 step of sgd
+        sgd_op.run()
+        # Validate updated params
+        self.assertAllCloseAccordingToType([[0, 1]],
+                                           self.evaluate(var0),
+                                           atol=0.01)
+
+  def testProximalAdagradWithL1(self):
+    # ProximalAdagradOptimizer is supported only in V1.
+    with ops.Graph().as_default(), self.cached_session():
+      var0 = variables.Variable([1.0, 2.0])
+      var1 = variables.Variable([4.0, 3.0])
+      grads0 = constant_op.constant([0.1, 0.2])
+      grads1 = constant_op.constant([0.01, 0.02])
+
+      opt = proximal_adagrad.ProximalAdagradOptimizer(
+          3.0,
+          initial_accumulator_value=0.1,
+          l1_regularization_strength=0.001,
+          l2_regularization_strength=0.0)
+      update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+      self.evaluate(variables.global_variables_initializer())
+
+      v0_val, v1_val = self.evaluate([var0, var1])
+      self.assertAllClose([1.0, 2.0], v0_val)
+      self.assertAllClose([4.0, 3.0], v1_val)
+
+      # Run 10 steps Proximal Adagrad
+      for _ in range(10):
+        update.run()
+      v0_val, v1_val = self.evaluate([var0, var1])
+      self.assertAllClose(np.array([-6.663634, -9.190331]), v0_val)
+      self.assertAllClose(np.array([2.959304, 1.029232]), v1_val)
+
+  def testProximalAdagradWithL1_L2(self):
+    # ProximalAdagradOptimizer is supported only in V1.
+    with ops.Graph().as_default(), self.cached_session():
+      var0 = variables.Variable([1.0, 2.0])
+      var1 = variables.Variable([4.0, 3.0])
+      grads0 = constant_op.constant([0.1, 0.2])
+      grads1 = constant_op.constant([0.01, 0.02])
+
+      opt = proximal_adagrad.ProximalAdagradOptimizer(
+          3.0,
+          initial_accumulator_value=0.1,
+          l1_regularization_strength=0.001,
+          l2_regularization_strength=2.0)
+      update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+      self.evaluate(variables.global_variables_initializer())
+
+      v0_val, v1_val = self.evaluate([var0, var1])
+      self.assertAllClose([1.0, 2.0], v0_val)
+      self.assertAllClose([4.0, 3.0], v1_val)
+
+      # Run 10 steps Proximal Adagrad.
+      for _ in range(10):
+        update.run()
+
+      v0_val, v1_val = self.evaluate([var0, var1])
+      self.assertAllClose(np.array([-0.0495, -0.0995]), v0_val)
+      self.assertAllClose(np.array([-0.0045, -0.0095]), v1_val)
+
+  def applyOptimizer(self, opt, steps=5, is_sparse=False):
+    if is_sparse:
+      var0 = variables.Variable([[1.0], [2.0]])
+      var1 = variables.Variable([[3.0], [4.0]])
+      grads0 = indexed_slices.IndexedSlices(
+          constant_op.constant(
+              [0.1], shape=[1, 1]),
+          constant_op.constant([0]),
+          constant_op.constant([2, 1]))
+      grads1 = indexed_slices.IndexedSlices(
+          constant_op.constant(
+              [0.02], shape=[1, 1]),
+          constant_op.constant([1]),
+          constant_op.constant([2, 1]))
+    else:
+      var0 = variables.Variable([1.0, 2.0])
+      var1 = variables.Variable([3.0, 4.0])
+      grads0 = constant_op.constant([0.1, 0.2])
+      grads1 = constant_op.constant([0.01, 0.02])
+
+    update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+    self.evaluate(variables.global_variables_initializer())
+
+    sess = ops.get_default_session()
+    v0_val, v1_val = self.evaluate([var0, var1])
+    if is_sparse:
+      self.assertAllClose([[1.0], [2.0]], v0_val)
+      self.assertAllClose([[3.0], [4.0]], v1_val)
+    else:
+      self.assertAllClose([1.0, 2.0], v0_val)
+      self.assertAllClose([3.0, 4.0], v1_val)
+
+    # Run ProximalAdagrad for a few steps
+    for _ in range(steps):
+      update.run()
+
+    v0_val, v1_val = self.evaluate([var0, var1])
+    return v0_val, v1_val
+
+  def testEquivAdagradwithoutRegularization(self):
+    # ProximalAdagradOptimizer is supported only in V1.
+    with ops.Graph().as_default(), self.cached_session():
+      val0, val1 = self.applyOptimizer(
+          proximal_adagrad.ProximalAdagradOptimizer(
+              3.0,
+              initial_accumulator_value=0.1,
+              l1_regularization_strength=0.0,
+              l2_regularization_strength=0.0))
+
+    with ops.Graph().as_default(), self.cached_session():
+      val2, val3 = self.applyOptimizer(
+          adagrad.AdagradOptimizer(
+              3.0, initial_accumulator_value=0.1))
+
+    self.assertAllClose(val0, val2)
+    self.assertAllClose(val1, val3)
+
+  def testEquivSparseAdagradwithoutRegularization(self):
+    # ProximalAdagradOptimizer is supported only in V1.
+    with ops.Graph().as_default(), self.cached_session():
+      val0, val1 = self.applyOptimizer(
+          proximal_adagrad.ProximalAdagradOptimizer(
+              3.0,
+              initial_accumulator_value=0.1,
+              l1_regularization_strength=0.0,
+              l2_regularization_strength=0.0),
+          is_sparse=True)
+
+    with ops.Graph().as_default(), self.cached_session():
+      val2, val3 = self.applyOptimizer(
+          adagrad.AdagradOptimizer(
+              3.0, initial_accumulator_value=0.1),
+          is_sparse=True)
+
+    self.assertAllClose(val0, val2)
+    self.assertAllClose(val1, val3)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/proximal_gradient_descent.py

@@ -0,0 +1,104 @@
+# 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.
+# ==============================================================================
+
+"""ProximalGradientDescent for TensorFlow."""
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import gen_training_ops
+# pylint: disable=unused-import
+from tensorflow.python.ops import math_ops
+# pylint: enable=unused-import
+from tensorflow.python.training import optimizer
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(v1=["train.ProximalGradientDescentOptimizer"])
+class ProximalGradientDescentOptimizer(optimizer.Optimizer):
+  # pylint: disable=line-too-long
+  """Optimizer that implements the proximal gradient descent algorithm.
+
+  References:
+    Efficient Learning using Forward-Backward Splitting:
+      [Duchi et al., 2009](http://papers.nips.cc/paper/3793-efficient-learning-using-forward-backward-splitting)
+      ([pdf](http://papers.nips.cc/paper/3793-efficient-learning-using-forward-backward-splitting.pdf))
+  """
+
+  def __init__(self, learning_rate, l1_regularization_strength=0.0,
+               l2_regularization_strength=0.0, use_locking=False,
+               name="ProximalGradientDescent"):
+    """Construct a new proximal gradient descent optimizer.
+
+    Args:
+      learning_rate: A Tensor or a floating point value.  The learning
+        rate to use.
+      l1_regularization_strength: A float value, must be greater than or
+        equal to zero.
+      l2_regularization_strength: A float value, must be greater than or
+        equal to zero.
+      use_locking: If True use locks for update operations.
+      name: Optional name prefix for the operations created when applying
+        gradients. Defaults to "GradientDescent".
+    """
+    super(ProximalGradientDescentOptimizer, self).__init__(use_locking, name)
+    self._learning_rate = learning_rate
+    self._l1_regularization_strength = l1_regularization_strength
+    self._l2_regularization_strength = l2_regularization_strength
+    self._l1_regularization_strength_tensor = None
+    self._l2_regularization_strength_tensor = None
+
+  def _apply_dense(self, grad, var):
+    return gen_training_ops.apply_proximal_gradient_descent(
+        var,
+        self._learning_rate_tensor,
+        self._l1_regularization_strength_tensor,
+        self._l2_regularization_strength_tensor,
+        grad,
+        use_locking=self._use_locking).op
+
+  def _resource_apply_dense(self, grad, var):
+    return gen_training_ops.resource_apply_proximal_gradient_descent(
+        var.handle,
+        self._learning_rate_tensor,
+        self._l1_regularization_strength_tensor,
+        self._l2_regularization_strength_tensor,
+        grad,
+        use_locking=self._use_locking)
+
+  def _apply_sparse(self, grad, var):
+    return gen_training_ops.sparse_apply_proximal_gradient_descent(
+        var,
+        self._learning_rate_tensor,
+        self._l1_regularization_strength_tensor,
+        self._l2_regularization_strength_tensor,
+        grad.values,
+        grad.indices,
+        use_locking=self._use_locking).op
+
+  def _resource_apply_sparse(self, grad, var, indices):
+    return gen_training_ops.resource_sparse_apply_proximal_gradient_descent(
+        var.handle,
+        math_ops.cast(self._learning_rate_tensor, grad.dtype),
+        math_ops.cast(self._l1_regularization_strength_tensor, grad.dtype),
+        math_ops.cast(self._l2_regularization_strength_tensor, grad.dtype),
+        grad,
+        indices,
+        use_locking=self._use_locking)
+
+  def _prepare(self):
+    self._learning_rate_tensor = ops.convert_to_tensor(self._learning_rate,
+                                                       name="learning_rate")
+    self._l1_regularization_strength_tensor = ops.convert_to_tensor(
+        self._l1_regularization_strength, name="l1_regularization_strength")
+    self._l2_regularization_strength_tensor = ops.convert_to_tensor(
+        self._l2_regularization_strength, name="l2_regularization_strength")

tensorflow/python/training/proximal_gradient_descent_test.py

@@ -0,0 +1,204 @@
+# 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 tests for Proximal Gradient Descent operations."""
+
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import embedding_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import gradient_descent
+from tensorflow.python.training import proximal_gradient_descent
+
+
+class ProximalGradientDescentOptimizerTest(test.TestCase):
+
+  def doTestProximalGradientDescentwithoutRegularization(
+      self, use_resource=False):
+    with ops.Graph().as_default(), self.cached_session():
+      if use_resource:
+        var0 = resource_variable_ops.ResourceVariable([0.0, 0.0])
+        var1 = resource_variable_ops.ResourceVariable([0.0, 0.0])
+      else:
+        var0 = variables.Variable([0.0, 0.0])
+        var1 = variables.Variable([0.0, 0.0])
+      grads0 = constant_op.constant([0.1, 0.2])
+      grads1 = constant_op.constant([0.01, 0.02])
+      opt = proximal_gradient_descent.ProximalGradientDescentOptimizer(
+          3.0, l1_regularization_strength=0.0, l2_regularization_strength=0.0)
+      update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+      self.evaluate(variables.global_variables_initializer())
+
+      v0_val, v1_val = self.evaluate([var0, var1])
+      self.assertAllClose([0.0, 0.0], v0_val)
+      self.assertAllClose([0.0, 0.0], v1_val)
+
+      # Run 3 steps Proximal Gradient Descent.
+      for _ in range(3):
+        update.run()
+
+      v0_val, v1_val = self.evaluate([var0, var1])
+      self.assertAllClose(np.array([-0.9, -1.8]), v0_val)
+      self.assertAllClose(np.array([-0.09, -0.18]), v1_val)
+
+  def testProximalGradientDescentwithoutRegularization(self):
+    self.doTestProximalGradientDescentwithoutRegularization(use_resource=False)
+
+  def testResourceProximalGradientDescentwithoutRegularization(self):
+    self.doTestProximalGradientDescentwithoutRegularization(use_resource=True)
+
+  def testProximalGradientDescentwithoutRegularization2(self):
+    with ops.Graph().as_default(), self.cached_session():
+      var0 = variables.Variable([1.0, 2.0])
+      var1 = variables.Variable([4.0, 3.0])
+      grads0 = constant_op.constant([0.1, 0.2])
+      grads1 = constant_op.constant([0.01, 0.02])
+
+      opt = proximal_gradient_descent.ProximalGradientDescentOptimizer(
+          3.0, l1_regularization_strength=0.0, l2_regularization_strength=0.0)
+      update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+      self.evaluate(variables.global_variables_initializer())
+
+      v0_val, v1_val = self.evaluate([var0, var1])
+      self.assertAllClose([1.0, 2.0], v0_val)
+      self.assertAllClose([4.0, 3.0], v1_val)
+
+      # Run 3 steps Proximal Gradient Descent
+      for _ in range(3):
+        update.run()
+
+      v0_val, v1_val = self.evaluate([var0, var1])
+      self.assertAllClose(np.array([0.1, 0.2]), v0_val)
+      self.assertAllClose(np.array([3.91, 2.82]), v1_val)
+
+  def testMinimizeSparseResourceVariable(self):
+    for dtype in [dtypes.float32, dtypes.float64]:
+      with ops.Graph().as_default(), self.cached_session():
+        var0 = resource_variable_ops.ResourceVariable([[1.0, 2.0]], dtype=dtype)
+        x = constant_op.constant([[4.0], [5.0]], dtype=dtype)
+        pred = math_ops.matmul(embedding_ops.embedding_lookup([var0], [0]), x)
+        loss = pred * pred
+        sgd_op = proximal_gradient_descent.ProximalGradientDescentOptimizer(
+            1.0).minimize(loss)
+        self.evaluate(variables.global_variables_initializer())
+        # Fetch params to validate initial values
+        self.assertAllCloseAccordingToType([[1.0, 2.0]], self.evaluate(var0))
+        # Run 1 step of sgd
+        sgd_op.run()
+        # Validate updated params
+        self.assertAllCloseAccordingToType([[-111, -138]],
+                                           self.evaluate(var0),
+                                           atol=0.01)
+
+  def testProximalGradientDescentWithL1_L2(self):
+    with ops.Graph().as_default(), self.cached_session():
+      var0 = variables.Variable([1.0, 2.0])
+      var1 = variables.Variable([4.0, 3.0])
+      grads0 = constant_op.constant([0.1, 0.2])
+      grads1 = constant_op.constant([0.01, 0.02])
+
+      opt = proximal_gradient_descent.ProximalGradientDescentOptimizer(
+          3.0, l1_regularization_strength=0.001, l2_regularization_strength=2.0)
+      update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+      self.evaluate(variables.global_variables_initializer())
+
+      v0_val, v1_val = self.evaluate([var0, var1])
+      self.assertAllClose([1.0, 2.0], v0_val)
+      self.assertAllClose([4.0, 3.0], v1_val)
+
+      # Run 10 steps Proximal Gradient Descent
+      for _ in range(10):
+        update.run()
+
+      v0_val, v1_val = self.evaluate([var0, var1])
+      self.assertAllClose(np.array([-0.0495, -0.0995]), v0_val)
+      self.assertAllClose(np.array([-0.0045, -0.0095]), v1_val)
+
+  def applyOptimizer(self, opt, steps=5, is_sparse=False):
+    if is_sparse:
+      var0 = variables.Variable([[1.0], [2.0]])
+      var1 = variables.Variable([[3.0], [4.0]])
+      grads0 = indexed_slices.IndexedSlices(
+          constant_op.constant(
+              [0.1], shape=[1, 1]),
+          constant_op.constant([0]),
+          constant_op.constant([2, 1]))
+      grads1 = indexed_slices.IndexedSlices(
+          constant_op.constant(
+              [0.02], shape=[1, 1]),
+          constant_op.constant([1]),
+          constant_op.constant([2, 1]))
+    else:
+      var0 = variables.Variable([1.0, 2.0])
+      var1 = variables.Variable([3.0, 4.0])
+      grads0 = constant_op.constant([0.1, 0.2])
+      grads1 = constant_op.constant([0.01, 0.02])
+
+    update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+    self.evaluate(variables.global_variables_initializer())
+
+    v0_val, v1_val = self.evaluate([var0, var1])
+    if is_sparse:
+      self.assertAllClose([[1.0], [2.0]], v0_val)
+      self.assertAllClose([[3.0], [4.0]], v1_val)
+    else:
+      self.assertAllClose([1.0, 2.0], v0_val)
+      self.assertAllClose([3.0, 4.0], v1_val)
+
+    # Run ProximalAdagrad for a few steps
+    for _ in range(steps):
+      update.run()
+
+    v0_val, v1_val = self.evaluate([var0, var1])
+    return v0_val, v1_val
+
+  def testEquivSparseGradientDescentwithoutRegularization(self):
+    with ops.Graph().as_default(), self.cached_session():
+      val0, val1 = self.applyOptimizer(
+          proximal_gradient_descent.ProximalGradientDescentOptimizer(
+              3.0,
+              l1_regularization_strength=0.0,
+              l2_regularization_strength=0.0),
+          is_sparse=True)
+
+      val2, val3 = self.applyOptimizer(
+          gradient_descent.GradientDescentOptimizer(3.0), is_sparse=True)
+
+    self.assertAllClose(val0, val2)
+    self.assertAllClose(val1, val3)
+
+  def testEquivGradientDescentwithoutRegularization(self):
+    with ops.Graph().as_default(), self.cached_session():
+      val0, val1 = self.applyOptimizer(
+          proximal_gradient_descent.ProximalGradientDescentOptimizer(
+              3.0,
+              l1_regularization_strength=0.0,
+              l2_regularization_strength=0.0))
+
+      val2, val3 = self.applyOptimizer(
+          gradient_descent.GradientDescentOptimizer(3.0))
+
+    self.assertAllClose(val0, val2)
+    self.assertAllClose(val1, val3)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/py_checkpoint_reader.py

@@ -0,0 +1,96 @@
+# 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.
+# ==============================================================================
+"""Extending CheckpointReader for TensorFlow."""
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.util import compat
+from tensorflow.python.util._pywrap_checkpoint_reader import CheckpointReader
+from tensorflow.python.util.tf_export import tf_export
+
+
+def error_translator(e):
+  """Translate the tensor_slice_reader.cc errors."""
+  # TODO(b/143319754): Remove the RuntimeError casting logic once we resolve the
+  # issue with throwing python exceptions from C++.
+  error_message = str(e)
+  if 'not found in checkpoint' in error_message or (
+      'Failed to find any '
+      'matching files for') in error_message:
+    raise errors_impl.NotFoundError(None, None, error_message)
+  elif 'Sliced checkpoints are not supported' in error_message or (
+      'Data type '
+      'not '
+      'supported') in error_message:
+    raise errors_impl.UnimplementedError(None, None, error_message)
+  elif 'Failed to get matching files on' in error_message:
+    raise errors_impl.InvalidArgumentError(None, None, error_message)
+  elif 'Unable to open table file' in error_message:
+    raise errors_impl.DataLossError(None, None, error_message)
+  elif 'Failed to find the saved tensor slices' in error_message or (
+      'not convertible to numpy dtype' in error_message):
+    raise errors_impl.InternalError(None, None, error_message)
+  else:
+    raise errors_impl.OpError(None, None, error_message, errors_impl.UNKNOWN)
+
+
+def get_variable_to_dtype_map(self):
+  return {
+      name: dtypes.DType(type_enum)
+      for name, type_enum in self._GetVariableToDataTypeMap().items()  # pylint: disable=protected-access
+  }
+
+CheckpointReader.get_variable_to_dtype_map = get_variable_to_dtype_map
+
+
+def has_tensor(self, tensor_str):
+  return self._HasTensor(compat.as_bytes(tensor_str))  # pylint: disable=protected-access
+
+CheckpointReader.has_tensor = has_tensor
+
+
+def get_tensor(self, tensor_str):
+  """Get the tensor from the Checkpoint object."""
+  try:
+    return CheckpointReader.CheckpointReader_GetTensor(
+        self, compat.as_bytes(tensor_str))
+  # TODO(b/143319754): Remove the RuntimeError casting logic once we resolve the
+  # issue with throwing python exceptions from C++.
+  except RuntimeError as e:
+    error_translator(e)
+
+
+CheckpointReader.get_tensor = get_tensor
+
+
+# Disable invalid name to keep backwards compatibility with that function.
+# It was previously exported from py_checkpoint_reader.i which did not conform
+# to pylint checks.
+# pylint: disable=invalid-name
+@tf_export(v1=['train.NewCheckpointReader'])
+def NewCheckpointReader(filepattern):
+  """A function that returns a CheckPointReader.
+
+  Args:
+    filepattern: The filename.
+
+  Returns:
+    A CheckpointReader object.
+  """
+  try:
+    return CheckpointReader(compat.as_bytes(filepattern))
+  # TODO(b/143319754): Remove the RuntimeError casting logic once we resolve the
+  # issue with throwing python exceptions from C++.
+  except RuntimeError as e:
+    error_translator(e)

tensorflow/python/training/quantize_training.py

@@ -0,0 +1,46 @@
+# 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.
+# ==============================================================================
+"""Quantize training for TensorFlow."""
+from tensorflow.core.framework import graph_pb2
+from tensorflow.python._pywrap_quantize_training import DoQuantizeTrainingOnGraphDefHelper
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+# Migrated this python code from deprecated quantize_training.i
+@deprecation.deprecated(
+    None,
+    "GraphDef quantized training rewriter is deprecated in the long term.")
+@tf_export(v1=["train.do_quantize_training_on_graphdef"])
+def do_quantize_training_on_graphdef(input_graph, num_bits):
+  """A general quantization scheme is being developed in `tf.contrib.quantize`.
+
+  Consider using that instead, though since it is in the tf.contrib namespace,
+  it is not subject to backward compatibility guarantees.
+
+  Args:
+    input_graph: A `GraphDef`.
+    num_bits: The number of bits for quantize training.
+
+  Returns:
+    The graph with quantize training done.
+  """
+
+  graph = graph_pb2.GraphDef()
+  result_graph_string = DoQuantizeTrainingOnGraphDefHelper(
+      input_graph.SerializeToString(), num_bits)
+
+  graph.ParseFromString(result_graph_string)
+  return graph

tensorflow/python/training/quantize_training_test.py

@@ -0,0 +1,99 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Tests for the SWIG-wrapped quantize training rewriting."""
+
+import os
+
+from tensorflow.python.client import session
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import importer
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import quantize_training
+from tensorflow.python.training import saver as saver_module
+
+
+class PywrapQuantizeTrainingTest(test.TestCase):
+
+  # Mainly to verify the python interface is working.
+  # More tests for this function can be found in the related c++ tests.
+  def testQuantizeTraining(self):
+    with session.Session() as sess:
+      a = constant_op.constant(6.0, shape=[1, 1])
+      b = constant_op.constant(7.0, shape=[1, 1])
+      c = math_ops.matmul(a, b, name='matmul')
+
+      self.assertEqual(c.eval(), 42.0)
+      self.assertEqual(len(sess.graph_def.node), 3)
+
+      result = quantize_training.do_quantize_training_on_graphdef(
+          sess.graph_def, 8)
+
+      # We just want to guarantee that some rewrite happened.
+      self.assertGreater(len(result.node), 3)
+
+  # Test that save/restoring works for EMA variables generated in the
+  # quantized training rewrite.
+  @test_util.run_v1_only('The API is only expect to work with v1 session mode.')
+  def testQuantizedSaveRestore(self):
+    save_path = os.path.join(self.get_temp_dir(), 'quantized_save_restore')
+
+    g = ops.Graph()
+    with session.Session(graph=g) as sess:
+      a = constant_op.constant(6.0, shape=[1, 1], name='a')
+      b = variable_v1.VariableV1(
+          constant_op.constant(7.0, shape=[1, 1]), name='b')
+      c = math_ops.matmul(a, b, name='matmul')
+
+      init_op = variables.global_variables_initializer()
+
+      saver = saver_module.Saver({'b': b})
+
+      result = quantize_training.do_quantize_training_on_graphdef(
+          sess.graph_def, 8)
+
+    with ops.Graph().as_default() as g, session.Session(graph=g) as sess:
+      _ = importer.import_graph_def(result, name='')
+
+      # Initialize the variable.
+      self.evaluate(g.get_operation_by_name(init_op.name))
+
+      # Run the graph for one step to assign values to the quantization min/max
+      # variables.
+      self.evaluate(g.get_tensor_by_name(c.name))
+
+      saver.save(sess, save_path)
+
+    with ops.Graph().as_default() as g, session.Session(graph=g) as sess:
+      _ = importer.import_graph_def(result, name='')
+
+      # When we restore the saved variabled, the quantization variables should
+      # be restored as well.
+      saver.restore(sess, save_path)
+      self.assertEqual(7.0, sess.run(g.get_tensor_by_name('b:0')))
+      self.assertEqual(6.0, sess.run(g.get_tensor_by_name('a/Min/Variable:0')))
+      self.assertEqual(6.0, sess.run(g.get_tensor_by_name('a/Max/Variable:0')))
+      self.assertEqual(7.0,
+                       sess.run(g.get_tensor_by_name('b/read/Min/Variable:0')))
+      self.assertEqual(7.0,
+                       sess.run(g.get_tensor_by_name('b/read/Max/Variable:0')))
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/training/quantize_training_wrapper.cc

@@ -0,0 +1,55 @@
+/* Copyright 2019 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.
+==============================================================================*/
+
+#include <Python.h>
+
+#include <string>
+
+#include "absl/status/status.h"
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "tensorflow/core/common_runtime/quantize_training.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/python/lib/core/pybind11_lib.h"
+#include "tensorflow/python/lib/core/pybind11_status.h"
+
+namespace py = pybind11;
+
+namespace tensorflow {
+static PyObject* DoQuantizeTrainingOnGraphDefHelper(
+    const std::string& input_graph, int num_bits) {
+  std::string result;
+  // TODO(suharshs): Make the QuantizeAndDequantizeV2 configurable.
+  tensorflow::MaybeRaiseFromStatus(
+      tensorflow::DoQuantizeTrainingOnSerializedGraphDef(
+          input_graph, num_bits, "QuantizeAndDequantizeV2", &result));
+
+  PyObject* py_str = PyBytes_FromStringAndSize(result.data(), result.size());
+  if (!py_str) {
+    tensorflow::MaybeRaiseFromStatus(absl::InternalError(
+        "Failed to generate serialized string of the rewritten graph."));
+  }
+  return py_str;
+}
+}  // namespace tensorflow
+
+PYBIND11_MODULE(_pywrap_quantize_training, m) {
+  m.def("DoQuantizeTrainingOnGraphDefHelper",
+        [](const py::object input_graph, int num_bits) {
+          return tensorflow::PyoOrThrow(
+              tensorflow::DoQuantizeTrainingOnGraphDefHelper(
+                  input_graph.cast<std::string>(), num_bits));
+        });
+};

tensorflow/python/training/queue_runner.py

@@ -0,0 +1,20 @@
+# 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.
+# ==============================================================================
+
+"""Create threads to run multiple enqueue ops."""
+# go/tf-wildcard-import
+# pylint: disable=wildcard-import
+from tensorflow.python.training.queue_runner_impl import *
+# pylint: enable=wildcard-import

tensorflow/python/training/queue_runner_impl.py

@@ -0,0 +1,490 @@
+# 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.
+# ==============================================================================
+
+"""Create threads to run multiple enqueue ops."""
+import threading
+import weakref
+
+from tensorflow.core.protobuf import queue_runner_pb2
+from tensorflow.python.client import session
+from tensorflow.python.eager import context
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+_DEPRECATION_INSTRUCTION = (
+    "To construct input pipelines, use the `tf.data` module.")
+
+
+@tf_export(v1=["train.queue_runner.QueueRunner", "train.QueueRunner"])
+class QueueRunner:
+  """Holds a list of enqueue operations for a queue, each to be run in a thread.
+
+  Queues are a convenient TensorFlow mechanism to compute tensors
+  asynchronously using multiple threads. For example in the canonical 'Input
+  Reader' setup one set of threads generates filenames in a queue; a second set
+  of threads read records from the files, processes them, and enqueues tensors
+  on a second queue; a third set of threads dequeues these input records to
+  construct batches and runs them through training operations.
+
+  There are several delicate issues when running multiple threads that way:
+  closing the queues in sequence as the input is exhausted, correctly catching
+  and reporting exceptions, etc.
+
+  The `QueueRunner`, combined with the `Coordinator`, helps handle these issues.
+
+  @compatibility(TF2)
+  QueueRunners are not compatible with eager execution. Instead, please
+  use [tf.data](https://www.tensorflow.org/guide/data) to get data into your
+  model.
+  @end_compatibility
+  """
+
+  @deprecation.deprecated(None, _DEPRECATION_INSTRUCTION)
+  def __init__(self, queue=None, enqueue_ops=None, close_op=None,
+               cancel_op=None, queue_closed_exception_types=None,
+               queue_runner_def=None, import_scope=None):
+    """Create a QueueRunner.
+
+    On construction the `QueueRunner` adds an op to close the queue.  That op
+    will be run if the enqueue ops raise exceptions.
+
+    When you later call the `create_threads()` method, the `QueueRunner` will
+    create one thread for each op in `enqueue_ops`.  Each thread will run its
+    enqueue op in parallel with the other threads.  The enqueue ops do not have
+    to all be the same op, but it is expected that they all enqueue tensors in
+    `queue`.
+
+    Args:
+      queue: A `Queue`.
+      enqueue_ops: List of enqueue ops to run in threads later.
+      close_op: Op to close the queue. Pending enqueue ops are preserved.
+      cancel_op: Op to close the queue and cancel pending enqueue ops.
+      queue_closed_exception_types: Optional tuple of Exception types that
+        indicate that the queue has been closed when raised during an enqueue
+        operation.  Defaults to `(tf.errors.OutOfRangeError,)`.  Another common
+        case includes `(tf.errors.OutOfRangeError, tf.errors.CancelledError)`,
+        when some of the enqueue ops may dequeue from other Queues.
+      queue_runner_def: Optional `QueueRunnerDef` protocol buffer. If specified,
+        recreates the QueueRunner from its contents. `queue_runner_def` and the
+        other arguments are mutually exclusive.
+      import_scope: Optional `string`. Name scope to add. Only used when
+        initializing from protocol buffer.
+
+    Raises:
+      ValueError: If both `queue_runner_def` and `queue` are both specified.
+      ValueError: If `queue` or `enqueue_ops` are not provided when not
+        restoring from `queue_runner_def`.
+      RuntimeError: If eager execution is enabled.
+    """
+    if context.executing_eagerly():
+      raise RuntimeError(
+          "QueueRunners are not supported when eager execution is enabled. "
+          "Instead, please use tf.data to get data into your model.")
+
+    if queue_runner_def:
+      if queue or enqueue_ops:
+        raise ValueError("queue_runner_def and queue are mutually exclusive.")
+      self._init_from_proto(queue_runner_def,
+                            import_scope=import_scope)
+    else:
+      self._init_from_args(
+          queue=queue, enqueue_ops=enqueue_ops,
+          close_op=close_op, cancel_op=cancel_op,
+          queue_closed_exception_types=queue_closed_exception_types)
+    # Protect the count of runs to wait for.
+    self._lock = threading.Lock()
+    # A map from a session object to the number of outstanding queue runner
+    # threads for that session.
+    self._runs_per_session = weakref.WeakKeyDictionary()
+    # List of exceptions raised by the running threads.
+    self._exceptions_raised = []
+
+  def _init_from_args(self, queue=None, enqueue_ops=None, close_op=None,
+                      cancel_op=None, queue_closed_exception_types=None):
+    """Create a QueueRunner from arguments.
+
+    Args:
+      queue: A `Queue`.
+      enqueue_ops: List of enqueue ops to run in threads later.
+      close_op: Op to close the queue. Pending enqueue ops are preserved.
+      cancel_op: Op to close the queue and cancel pending enqueue ops.
+      queue_closed_exception_types: Tuple of exception types, which indicate
+        the queue has been safely closed.
+
+    Raises:
+      ValueError: If `queue` or `enqueue_ops` are not provided when not
+        restoring from `queue_runner_def`.
+      TypeError: If `queue_closed_exception_types` is provided, but is not
+        a non-empty tuple of error types (subclasses of `tf.errors.OpError`).
+    """
+    if not queue or not enqueue_ops:
+      raise ValueError("Must provide queue and enqueue_ops.")
+    self._queue = queue
+    self._enqueue_ops = enqueue_ops
+    self._close_op = close_op
+    self._cancel_op = cancel_op
+    if queue_closed_exception_types is not None:
+      if (not isinstance(queue_closed_exception_types, tuple)
+          or not queue_closed_exception_types
+          or not all(issubclass(t, errors.OpError)
+                     for t in queue_closed_exception_types)):
+        raise TypeError(
+            "queue_closed_exception_types, when provided, "
+            "must be a tuple of tf.error types, but saw: %s"
+            % queue_closed_exception_types)
+    self._queue_closed_exception_types = queue_closed_exception_types
+    # Close when no more will be produced, but pending enqueues should be
+    # preserved.
+    if self._close_op is None:
+      self._close_op = self._queue.close()
+    # Close and cancel pending enqueues since there was an error and we want
+    # to unblock everything so we can cleanly exit.
+    if self._cancel_op is None:
+      self._cancel_op = self._queue.close(cancel_pending_enqueues=True)
+    if not self._queue_closed_exception_types:
+      self._queue_closed_exception_types = (errors.OutOfRangeError,)
+    else:
+      self._queue_closed_exception_types = tuple(
+          self._queue_closed_exception_types)
+
+  def _init_from_proto(self, queue_runner_def, import_scope=None):
+    """Create a QueueRunner from `QueueRunnerDef`.
+
+    Args:
+      queue_runner_def: Optional `QueueRunnerDef` protocol buffer.
+      import_scope: Optional `string`. Name scope to add.
+    """
+    assert isinstance(queue_runner_def, queue_runner_pb2.QueueRunnerDef)
+    g = ops.get_default_graph()
+    self._queue = g.as_graph_element(
+        ops.prepend_name_scope(queue_runner_def.queue_name, import_scope))
+    self._enqueue_ops = [g.as_graph_element(
+        ops.prepend_name_scope(op, import_scope))
+                         for op in queue_runner_def.enqueue_op_name]
+    self._close_op = g.as_graph_element(ops.prepend_name_scope(
+        queue_runner_def.close_op_name, import_scope))
+    self._cancel_op = g.as_graph_element(ops.prepend_name_scope(
+        queue_runner_def.cancel_op_name, import_scope))
+    self._queue_closed_exception_types = tuple(
+        errors.exception_type_from_error_code(code)
+        for code in queue_runner_def.queue_closed_exception_types)
+    # Legacy support for old QueueRunnerDefs created before this field
+    # was added.
+    if not self._queue_closed_exception_types:
+      self._queue_closed_exception_types = (errors.OutOfRangeError,)
+
+  @property
+  def queue(self):
+    return self._queue
+
+  @property
+  def enqueue_ops(self):
+    return self._enqueue_ops
+
+  @property
+  def close_op(self):
+    return self._close_op
+
+  @property
+  def cancel_op(self):
+    return self._cancel_op
+
+  @property
+  def queue_closed_exception_types(self):
+    return self._queue_closed_exception_types
+
+  @property
+  def exceptions_raised(self):
+    """Exceptions raised but not handled by the `QueueRunner` threads.
+
+    Exceptions raised in queue runner threads are handled in one of two ways
+    depending on whether or not a `Coordinator` was passed to
+    `create_threads()`:
+
+    * With a `Coordinator`, exceptions are reported to the coordinator and
+      forgotten by the `QueueRunner`.
+    * Without a `Coordinator`, exceptions are captured by the `QueueRunner` and
+      made available in this `exceptions_raised` property.
+
+    Returns:
+      A list of Python `Exception` objects.  The list is empty if no exception
+      was captured.  (No exceptions are captured when using a Coordinator.)
+    """
+    return self._exceptions_raised
+
+  @property
+  def name(self):
+    """The string name of the underlying Queue."""
+    return self._queue.name
+
+  # pylint: disable=broad-except
+  def _run(self, sess, enqueue_op, coord=None):
+    """Execute the enqueue op in a loop, close the queue in case of error.
+
+    Args:
+      sess: A Session.
+      enqueue_op: The Operation to run.
+      coord: Optional Coordinator object for reporting errors and checking
+        for stop conditions.
+    """
+    decremented = False
+    try:
+      # Make a cached callable from the `enqueue_op` to decrease the
+      # Python overhead in the queue-runner loop.
+      enqueue_callable = sess.make_callable(enqueue_op)
+      while True:
+        if coord and coord.should_stop():
+          break
+        try:
+          enqueue_callable()
+        except self._queue_closed_exception_types:  # pylint: disable=catching-non-exception
+          # This exception indicates that a queue was closed.
+          with self._lock:
+            self._runs_per_session[sess] -= 1
+            decremented = True
+            if self._runs_per_session[sess] == 0:
+              try:
+                sess.run(self._close_op)
+              except Exception as e:
+                # Intentionally ignore errors from close_op.
+                logging.vlog(1, "Ignored exception: %s", str(e))
+            return
+    except Exception as e:
+      # This catches all other exceptions.
+      if coord:
+        coord.request_stop(e)
+      else:
+        logging.error("Exception in QueueRunner: %s", str(e))
+        with self._lock:
+          self._exceptions_raised.append(e)
+        raise
+    finally:
+      # Make sure we account for all terminations: normal or errors.
+      if not decremented:
+        with self._lock:
+          self._runs_per_session[sess] -= 1
+
+  def _close_on_stop(self, sess, cancel_op, coord):
+    """Close the queue when the Coordinator requests stop.
+
+    Args:
+      sess: A Session.
+      cancel_op: The Operation to run.
+      coord: Coordinator.
+    """
+    coord.wait_for_stop()
+    try:
+      sess.run(cancel_op)
+    except Exception as e:
+      # Intentionally ignore errors from cancel_op.
+      logging.vlog(1, "Ignored exception: %s", str(e))
+  # pylint: enable=broad-except
+
+  def create_threads(self, sess, coord=None, daemon=False, start=False):
+    """Create threads to run the enqueue ops for the given session.
+
+    This method requires a session in which the graph was launched.  It creates
+    a list of threads, optionally starting them.  There is one thread for each
+    op passed in `enqueue_ops`.
+
+    The `coord` argument is an optional coordinator that the threads will use
+    to terminate together and report exceptions.  If a coordinator is given,
+    this method starts an additional thread to close the queue when the
+    coordinator requests a stop.
+
+    If previously created threads for the given session are still running, no
+    new threads will be created.
+
+    Args:
+      sess: A `Session`.
+      coord: Optional `Coordinator` object for reporting errors and checking
+        stop conditions.
+      daemon: Boolean.  If `True` make the threads daemon threads.
+      start: Boolean.  If `True` starts the threads.  If `False` the
+        caller must call the `start()` method of the returned threads.
+
+    Returns:
+      A list of threads.
+    """
+    with self._lock:
+      try:
+        if self._runs_per_session[sess] > 0:
+          # Already started: no new threads to return.
+          return []
+      except KeyError:
+        # We haven't seen this session yet.
+        pass
+      self._runs_per_session[sess] = len(self._enqueue_ops)
+      self._exceptions_raised = []
+
+    ret_threads = []
+    for op in self._enqueue_ops:
+      name = "QueueRunnerThread-{}-{}".format(self.name, op.name)
+      ret_threads.append(threading.Thread(target=self._run,
+                                          args=(sess, op, coord),
+                                          name=name))
+    if coord:
+      name = "QueueRunnerThread-{}-close_on_stop".format(self.name)
+      ret_threads.append(threading.Thread(target=self._close_on_stop,
+                                          args=(sess, self._cancel_op, coord),
+                                          name=name))
+    for t in ret_threads:
+      if coord:
+        coord.register_thread(t)
+      if daemon:
+        t.daemon = True
+      if start:
+        t.start()
+    return ret_threads
+
+  def to_proto(self, export_scope=None):
+    """Converts this `QueueRunner` to a `QueueRunnerDef` protocol buffer.
+
+    Args:
+      export_scope: Optional `string`. Name scope to remove.
+
+    Returns:
+      A `QueueRunnerDef` protocol buffer, or `None` if the `Variable` is not in
+      the specified name scope.
+    """
+    if (export_scope is None or
+        self.queue.name.startswith(export_scope)):
+      queue_runner_def = queue_runner_pb2.QueueRunnerDef()
+      queue_runner_def.queue_name = ops.strip_name_scope(
+          self.queue.name, export_scope)
+      for enqueue_op in self.enqueue_ops:
+        queue_runner_def.enqueue_op_name.append(
+            ops.strip_name_scope(enqueue_op.name, export_scope))
+      queue_runner_def.close_op_name = ops.strip_name_scope(
+          self.close_op.name, export_scope)
+      queue_runner_def.cancel_op_name = ops.strip_name_scope(
+          self.cancel_op.name, export_scope)
+      queue_runner_def.queue_closed_exception_types.extend([
+          errors.error_code_from_exception_type(cls)
+          for cls in self._queue_closed_exception_types])
+      return queue_runner_def
+    else:
+      return None
+
+  @staticmethod
+  def from_proto(queue_runner_def, import_scope=None):
+    """Returns a `QueueRunner` object created from `queue_runner_def`."""
+    return QueueRunner(queue_runner_def=queue_runner_def,
+                       import_scope=import_scope)
+
+
+@tf_export(v1=["train.queue_runner.add_queue_runner", "train.add_queue_runner"])
+@deprecation.deprecated(None, _DEPRECATION_INSTRUCTION)
+def add_queue_runner(qr, collection=ops.GraphKeys.QUEUE_RUNNERS):
+  """Adds a `QueueRunner` to a collection in the graph.
+
+  When building a complex model that uses many queues it is often difficult to
+  gather all the queue runners that need to be run.  This convenience function
+  allows you to add a queue runner to a well known collection in the graph.
+
+  The companion method `start_queue_runners()` can be used to start threads for
+  all the collected queue runners.
+
+  @compatibility(TF2)
+  QueueRunners are not compatible with eager execution. Instead, please
+  use [tf.data](https://www.tensorflow.org/guide/data) to get data into your
+  model.
+  @end_compatibility
+
+  Args:
+    qr: A `QueueRunner`.
+    collection: A `GraphKey` specifying the graph collection to add
+      the queue runner to.  Defaults to `GraphKeys.QUEUE_RUNNERS`.
+  """
+  ops.add_to_collection(collection, qr)
+
+
+@tf_export(v1=["train.queue_runner.start_queue_runners",
+               "train.start_queue_runners"])
+@deprecation.deprecated(None, _DEPRECATION_INSTRUCTION)
+def start_queue_runners(sess=None, coord=None, daemon=True, start=True,
+                        collection=ops.GraphKeys.QUEUE_RUNNERS):
+  """Starts all queue runners collected in the graph.
+
+  This is a companion method to `add_queue_runner()`.  It just starts
+  threads for all queue runners collected in the graph.  It returns
+  the list of all threads.
+
+  @compatibility(TF2)
+  QueueRunners are not compatible with eager execution. Instead, please
+  use [tf.data](https://www.tensorflow.org/guide/data) to get data into your
+  model.
+  @end_compatibility
+
+  Args:
+    sess: `Session` used to run the queue ops.  Defaults to the
+      default session.
+    coord: Optional `Coordinator` for coordinating the started threads.
+    daemon: Whether the threads should be marked as `daemons`, meaning
+      they don't block program exit.
+    start: Set to `False` to only create the threads, not start them.
+    collection: A `GraphKey` specifying the graph collection to
+      get the queue runners from.  Defaults to `GraphKeys.QUEUE_RUNNERS`.
+
+  Raises:
+    ValueError: if `sess` is None and there isn't any default session.
+    TypeError: if `sess` is not a `tf.compat.v1.Session` object.
+
+  Returns:
+    A list of threads.
+
+  Raises:
+    RuntimeError: If called with eager execution enabled.
+    ValueError: If called without a default `tf.compat.v1.Session` registered.
+  """
+  if context.executing_eagerly():
+    raise RuntimeError("Queues are not compatible with eager execution.")
+  if sess is None:
+    sess = ops.get_default_session()
+    if not sess:
+      raise ValueError("Cannot start queue runners: No default session is "
+                       "registered. Use `with sess.as_default()` or pass an "
+                       "explicit session to tf.start_queue_runners(sess=sess)")
+
+  if not isinstance(sess, session.SessionInterface):
+    # Following check is due to backward compatibility. (b/62061352)
+    if sess.__class__.__name__ in [
+        "MonitoredSession", "SingularMonitoredSession"]:
+      return []
+    raise TypeError("sess must be a `tf.Session` object. "
+                    "Given class: {}".format(sess.__class__))
+
+  queue_runners = ops.get_collection(collection)
+  if not queue_runners:
+    logging.warning(
+        "`tf.train.start_queue_runners()` was called when no queue runners "
+        "were defined. You can safely remove the call to this deprecated "
+        "function.")
+
+  with sess.graph.as_default():
+    threads = []
+    for qr in ops.get_collection(collection):
+      threads.extend(qr.create_threads(sess, coord=coord, daemon=daemon,
+                                       start=start))
+  return threads
+
+
+ops.register_proto_function(ops.GraphKeys.QUEUE_RUNNERS,
+                            proto_type=queue_runner_pb2.QueueRunnerDef,
+                            to_proto=QueueRunner.to_proto,
+                            from_proto=QueueRunner.from_proto)

tensorflow/python/training/queue_runner_test.py

@@ -0,0 +1,352 @@
+# 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.
+# ==============================================================================
+"""Tests for QueueRunner."""
+
+import collections
+import time
+
+from tensorflow.python.client import session
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import data_flow_ops
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import coordinator
+from tensorflow.python.training import monitored_session
+from tensorflow.python.training import queue_runner_impl
+
+
+_MockOp = collections.namedtuple("MockOp", ["name"])
+
+
+@test_util.run_v1_only("QueueRunner removed from v2")
+class QueueRunnerTest(test.TestCase):
+
+  def testBasic(self):
+    with self.cached_session() as sess:
+      # CountUpTo will raise OUT_OF_RANGE when it reaches the count.
+      zero64 = constant_op.constant(0, dtype=dtypes.int64)
+      var = variable_v1.VariableV1(zero64)
+      count_up_to = var.count_up_to(3)
+      queue = data_flow_ops.FIFOQueue(10, dtypes.float32)
+      self.evaluate(variables.global_variables_initializer())
+      qr = queue_runner_impl.QueueRunner(queue, [count_up_to])
+      threads = qr.create_threads(sess)
+      self.assertEqual(sorted(t.name for t in threads),
+                       ["QueueRunnerThread-fifo_queue-CountUpTo:0"])
+      for t in threads:
+        t.start()
+      for t in threads:
+        t.join()
+      self.assertEqual(0, len(qr.exceptions_raised))
+      # The variable should be 3.
+      self.assertEqual(3, self.evaluate(var))
+
+  def testTwoOps(self):
+    with self.cached_session() as sess:
+      # CountUpTo will raise OUT_OF_RANGE when it reaches the count.
+      zero64 = constant_op.constant(0, dtype=dtypes.int64)
+      var0 = variable_v1.VariableV1(zero64)
+      count_up_to_3 = var0.count_up_to(3)
+      var1 = variable_v1.VariableV1(zero64)
+      count_up_to_30 = var1.count_up_to(30)
+      queue = data_flow_ops.FIFOQueue(10, dtypes.float32)
+      qr = queue_runner_impl.QueueRunner(queue, [count_up_to_3, count_up_to_30])
+      threads = qr.create_threads(sess)
+      self.assertEqual(sorted(t.name for t in threads),
+                       ["QueueRunnerThread-fifo_queue-CountUpTo:0",
+                        "QueueRunnerThread-fifo_queue-CountUpTo_1:0"])
+      self.evaluate(variables.global_variables_initializer())
+      for t in threads:
+        t.start()
+      for t in threads:
+        t.join()
+      self.assertEqual(0, len(qr.exceptions_raised))
+      self.assertEqual(3, self.evaluate(var0))
+      self.assertEqual(30, self.evaluate(var1))
+
+  def testExceptionsCaptured(self):
+    with self.cached_session() as sess:
+      queue = data_flow_ops.FIFOQueue(10, dtypes.float32)
+      qr = queue_runner_impl.QueueRunner(queue, [_MockOp("i fail"),
+                                                 _MockOp("so fail")])
+      threads = qr.create_threads(sess)
+      self.evaluate(variables.global_variables_initializer())
+      for t in threads:
+        t.start()
+      for t in threads:
+        t.join()
+      exceptions = qr.exceptions_raised
+      self.assertEqual(2, len(exceptions))
+      self.assertTrue("Operation not in the graph" in str(exceptions[0]))
+      self.assertTrue("Operation not in the graph" in str(exceptions[1]))
+
+  def testRealDequeueEnqueue(self):
+    with self.cached_session() as sess:
+      q0 = data_flow_ops.FIFOQueue(3, dtypes.float32)
+      enqueue0 = q0.enqueue((10.0,))
+      close0 = q0.close()
+      q1 = data_flow_ops.FIFOQueue(30, dtypes.float32)
+      enqueue1 = q1.enqueue((q0.dequeue(),))
+      dequeue1 = q1.dequeue()
+      qr = queue_runner_impl.QueueRunner(q1, [enqueue1])
+      threads = qr.create_threads(sess)
+      for t in threads:
+        t.start()
+      # Enqueue 2 values, then close queue0.
+      enqueue0.run()
+      enqueue0.run()
+      close0.run()
+      # Wait for the queue runner to terminate.
+      for t in threads:
+        t.join()
+      # It should have terminated cleanly.
+      self.assertEqual(0, len(qr.exceptions_raised))
+      # The 2 values should be in queue1.
+      self.assertEqual(10.0, self.evaluate(dequeue1))
+      self.assertEqual(10.0, self.evaluate(dequeue1))
+      # And queue1 should now be closed.
+      with self.assertRaisesRegex(errors_impl.OutOfRangeError, "is closed"):
+        self.evaluate(dequeue1)
+
+  def testRespectCoordShouldStop(self):
+    with self.cached_session() as sess:
+      # CountUpTo will raise OUT_OF_RANGE when it reaches the count.
+      zero64 = constant_op.constant(0, dtype=dtypes.int64)
+      var = variable_v1.VariableV1(zero64)
+      count_up_to = var.count_up_to(3)
+      queue = data_flow_ops.FIFOQueue(10, dtypes.float32)
+      self.evaluate(variables.global_variables_initializer())
+      qr = queue_runner_impl.QueueRunner(queue, [count_up_to])
+      # As the coordinator to stop.  The queue runner should
+      # finish immediately.
+      coord = coordinator.Coordinator()
+      coord.request_stop()
+      threads = qr.create_threads(sess, coord)
+      self.assertEqual(sorted(t.name for t in threads),
+                       ["QueueRunnerThread-fifo_queue-CountUpTo:0",
+                        "QueueRunnerThread-fifo_queue-close_on_stop"])
+      for t in threads:
+        t.start()
+      coord.join()
+      self.assertEqual(0, len(qr.exceptions_raised))
+      # The variable should be 0.
+      self.assertEqual(0, self.evaluate(var))
+
+  def testRequestStopOnException(self):
+    with self.cached_session() as sess:
+      queue = data_flow_ops.FIFOQueue(10, dtypes.float32)
+      qr = queue_runner_impl.QueueRunner(queue, [_MockOp("not an op")])
+      coord = coordinator.Coordinator()
+      threads = qr.create_threads(sess, coord)
+      for t in threads:
+        t.start()
+      # The exception should be re-raised when joining.
+      with self.assertRaisesRegex(ValueError, "Operation not in the graph"):
+        coord.join()
+
+  def testGracePeriod(self):
+    with self.cached_session() as sess:
+      # The enqueue will quickly block.
+      queue = data_flow_ops.FIFOQueue(2, dtypes.float32)
+      enqueue = queue.enqueue((10.0,))
+      dequeue = queue.dequeue()
+      qr = queue_runner_impl.QueueRunner(queue, [enqueue])
+      coord = coordinator.Coordinator()
+      qr.create_threads(sess, coord, start=True)
+      # Dequeue one element and then request stop.
+      dequeue.op.run()
+      time.sleep(0.02)
+      coord.request_stop()
+      # We should be able to join because the RequestStop() will cause
+      # the queue to be closed and the enqueue to terminate.
+      coord.join(stop_grace_period_secs=1.0)
+
+  def testMultipleSessions(self):
+    with self.cached_session() as sess:
+      with session.Session() as other_sess:
+        zero64 = constant_op.constant(0, dtype=dtypes.int64)
+        var = variable_v1.VariableV1(zero64)
+        count_up_to = var.count_up_to(3)
+        queue = data_flow_ops.FIFOQueue(10, dtypes.float32)
+        self.evaluate(variables.global_variables_initializer())
+        coord = coordinator.Coordinator()
+        qr = queue_runner_impl.QueueRunner(queue, [count_up_to])
+        # NOTE that this test does not actually start the threads.
+        threads = qr.create_threads(sess, coord=coord)
+        other_threads = qr.create_threads(other_sess, coord=coord)
+        self.assertEqual(len(threads), len(other_threads))
+
+  def testIgnoreMultiStarts(self):
+    with self.cached_session() as sess:
+      # CountUpTo will raise OUT_OF_RANGE when it reaches the count.
+      zero64 = constant_op.constant(0, dtype=dtypes.int64)
+      var = variable_v1.VariableV1(zero64)
+      count_up_to = var.count_up_to(3)
+      queue = data_flow_ops.FIFOQueue(10, dtypes.float32)
+      self.evaluate(variables.global_variables_initializer())
+      coord = coordinator.Coordinator()
+      qr = queue_runner_impl.QueueRunner(queue, [count_up_to])
+      threads = []
+      # NOTE that this test does not actually start the threads.
+      threads.extend(qr.create_threads(sess, coord=coord))
+      new_threads = qr.create_threads(sess, coord=coord)
+      self.assertEqual([], new_threads)
+
+  def testThreads(self):
+    with self.cached_session() as sess:
+      # CountUpTo will raise OUT_OF_RANGE when it reaches the count.
+      zero64 = constant_op.constant(0, dtype=dtypes.int64)
+      var = variable_v1.VariableV1(zero64)
+      count_up_to = var.count_up_to(3)
+      queue = data_flow_ops.FIFOQueue(10, dtypes.float32)
+      self.evaluate(variables.global_variables_initializer())
+      qr = queue_runner_impl.QueueRunner(queue, [count_up_to,
+                                                 _MockOp("bad_op")])
+      threads = qr.create_threads(sess, start=True)
+      self.assertEqual(sorted(t.name for t in threads),
+                       ["QueueRunnerThread-fifo_queue-CountUpTo:0",
+                        "QueueRunnerThread-fifo_queue-bad_op"])
+      for t in threads:
+        t.join()
+      exceptions = qr.exceptions_raised
+      self.assertEqual(1, len(exceptions))
+      self.assertTrue("Operation not in the graph" in str(exceptions[0]))
+
+      threads = qr.create_threads(sess, start=True)
+      for t in threads:
+        t.join()
+      exceptions = qr.exceptions_raised
+      self.assertEqual(1, len(exceptions))
+      self.assertTrue("Operation not in the graph" in str(exceptions[0]))
+
+  def testName(self):
+    with ops.name_scope("scope"):
+      queue = data_flow_ops.FIFOQueue(10, dtypes.float32, name="queue")
+    qr = queue_runner_impl.QueueRunner(queue, [control_flow_ops.no_op()])
+    self.assertEqual("scope/queue", qr.name)
+    queue_runner_impl.add_queue_runner(qr)
+    self.assertEqual(
+        1, len(ops.get_collection(ops.GraphKeys.QUEUE_RUNNERS, "scope")))
+
+  def testStartQueueRunners(self):
+    # CountUpTo will raise OUT_OF_RANGE when it reaches the count.
+    zero64 = constant_op.constant(0, dtype=dtypes.int64)
+    var = variable_v1.VariableV1(zero64)
+    count_up_to = var.count_up_to(3)
+    queue = data_flow_ops.FIFOQueue(10, dtypes.float32)
+    init_op = variables.global_variables_initializer()
+    qr = queue_runner_impl.QueueRunner(queue, [count_up_to])
+    queue_runner_impl.add_queue_runner(qr)
+    with self.cached_session() as sess:
+      init_op.run()
+      threads = queue_runner_impl.start_queue_runners(sess)
+      for t in threads:
+        t.join()
+      self.assertEqual(0, len(qr.exceptions_raised))
+      # The variable should be 3.
+      self.assertEqual(3, self.evaluate(var))
+
+  def testStartQueueRunnersRaisesIfNotASession(self):
+    zero64 = constant_op.constant(0, dtype=dtypes.int64)
+    var = variable_v1.VariableV1(zero64)
+    count_up_to = var.count_up_to(3)
+    queue = data_flow_ops.FIFOQueue(10, dtypes.float32)
+    init_op = variables.global_variables_initializer()
+    qr = queue_runner_impl.QueueRunner(queue, [count_up_to])
+    queue_runner_impl.add_queue_runner(qr)
+    with self.cached_session():
+      init_op.run()
+      with self.assertRaisesRegex(TypeError, "tf.Session"):
+        queue_runner_impl.start_queue_runners("NotASession")
+
+  def testStartQueueRunnersIgnoresMonitoredSession(self):
+    zero64 = constant_op.constant(0, dtype=dtypes.int64)
+    var = variable_v1.VariableV1(zero64)
+    count_up_to = var.count_up_to(3)
+    queue = data_flow_ops.FIFOQueue(10, dtypes.float32)
+    init_op = variables.global_variables_initializer()
+    qr = queue_runner_impl.QueueRunner(queue, [count_up_to])
+    queue_runner_impl.add_queue_runner(qr)
+    with self.cached_session():
+      init_op.run()
+      threads = queue_runner_impl.start_queue_runners(
+          monitored_session.MonitoredSession())
+      self.assertFalse(threads)
+
+  def testStartQueueRunnersNonDefaultGraph(self):
+    # CountUpTo will raise OUT_OF_RANGE when it reaches the count.
+    graph = ops.Graph()
+    with graph.as_default():
+      zero64 = constant_op.constant(0, dtype=dtypes.int64)
+      var = variable_v1.VariableV1(zero64)
+      count_up_to = var.count_up_to(3)
+      queue = data_flow_ops.FIFOQueue(10, dtypes.float32)
+      init_op = variables.global_variables_initializer()
+      qr = queue_runner_impl.QueueRunner(queue, [count_up_to])
+      queue_runner_impl.add_queue_runner(qr)
+    with self.session(graph=graph) as sess:
+      init_op.run()
+      threads = queue_runner_impl.start_queue_runners(sess)
+      for t in threads:
+        t.join()
+      self.assertEqual(0, len(qr.exceptions_raised))
+      # The variable should be 3.
+      self.assertEqual(3, self.evaluate(var))
+
+  def testQueueRunnerSerializationRoundTrip(self):
+    graph = ops.Graph()
+    with graph.as_default():
+      queue = data_flow_ops.FIFOQueue(10, dtypes.float32, name="queue")
+      enqueue_op = control_flow_ops.no_op(name="enqueue")
+      close_op = control_flow_ops.no_op(name="close")
+      cancel_op = control_flow_ops.no_op(name="cancel")
+      qr0 = queue_runner_impl.QueueRunner(
+          queue, [enqueue_op],
+          close_op,
+          cancel_op,
+          queue_closed_exception_types=(errors_impl.OutOfRangeError,
+                                        errors_impl.CancelledError))
+      qr0_proto = queue_runner_impl.QueueRunner.to_proto(qr0)
+      qr0_recon = queue_runner_impl.QueueRunner.from_proto(qr0_proto)
+      self.assertEqual("queue", qr0_recon.queue.name)
+      self.assertEqual(1, len(qr0_recon.enqueue_ops))
+      self.assertEqual(enqueue_op, qr0_recon.enqueue_ops[0])
+      self.assertEqual(close_op, qr0_recon.close_op)
+      self.assertEqual(cancel_op, qr0_recon.cancel_op)
+      self.assertEqual(
+          (errors_impl.OutOfRangeError, errors_impl.CancelledError),
+          qr0_recon.queue_closed_exception_types)
+
+      # Assert we reconstruct an OutOfRangeError for QueueRunners
+      # created before QueueRunnerDef had a queue_closed_exception_types field.
+      del qr0_proto.queue_closed_exception_types[:]
+      qr0_legacy_recon = queue_runner_impl.QueueRunner.from_proto(qr0_proto)
+      self.assertEqual("queue", qr0_legacy_recon.queue.name)
+      self.assertEqual(1, len(qr0_legacy_recon.enqueue_ops))
+      self.assertEqual(enqueue_op, qr0_legacy_recon.enqueue_ops[0])
+      self.assertEqual(close_op, qr0_legacy_recon.close_op)
+      self.assertEqual(cancel_op, qr0_legacy_recon.cancel_op)
+      self.assertEqual((errors_impl.OutOfRangeError,),
+                       qr0_legacy_recon.queue_closed_exception_types)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/rmsprop.py

@@ -0,0 +1,323 @@
+# 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.
+# ==============================================================================
+"""One-line documentation for rmsprop module.
+
+rmsprop algorithm [tieleman2012rmsprop]
+
+A detailed description of rmsprop.
+
+- maintain a moving (discounted) average of the square of gradients
+- divide gradient by the root of this average
+
+mean_square = decay * mean_square{t-1} + (1-decay) * gradient ** 2
+mom = momentum * mom{t-1} + learning_rate * g_t / sqrt(mean_square + epsilon)
+delta = - mom
+
+This implementation of RMSProp uses plain momentum, not Nesterov momentum.
+
+The centered version additionally maintains a moving (discounted) average of the
+gradients, and uses that average to estimate the variance:
+
+mean_grad = decay * mean_grad{t-1} + (1-decay) * gradient
+mean_square = decay * mean_square{t-1} + (1-decay) * gradient ** 2
+mom = momentum * mom{t-1} + learning_rate * g_t /
+    sqrt(mean_square - mean_grad**2 + epsilon)
+delta = - mom
+"""
+
+from tensorflow.python.framework import ops
+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
+from tensorflow.python.training import optimizer
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(v1=["train.RMSPropOptimizer"])
+class RMSPropOptimizer(optimizer.Optimizer):
+  """Optimizer that implements the RMSProp algorithm (Tielemans et al.
+
+  2012).
+
+  References:
+    Coursera slide 29:
+    Hinton, 2012
+    ([pdf](http://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf))
+
+  @compatibility(TF2)
+  tf.compat.v1.train.RMSPropOptimizer is compatible with eager mode and
+  `tf.function`.
+  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.
+
+  To switch to native TF2 style, use [`tf.keras.optimizers.RMSprop`]
+  (https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/RMSprop)
+  instead. Please notice that due to the implementation differences,
+  `tf.keras.optimizers.RMSprop` and
+  `tf.compat.v1.train.RMSPropOptimizer` may have slight differences in
+  floating point numerics even though the formula used for the variable
+  updates still matches.
+
+  #### Structural mapping to native TF2
+
+  Before:
+
+  ```python
+  optimizer = tf.compat.v1.train.RMSPropOptimizer(
+    learning_rate=learning_rate,
+    decay=decay,
+    momentum=momentum,
+    epsilon=epsilon)
+  ```
+
+  After:
+
+  ```python
+  optimizer = tf.keras.optimizers.RMSprop(
+    learning_rate=learning_rate,
+    rho=decay,
+    momentum=momentum,
+    epsilon=epsilon)
+  ```
+
+  #### How to map arguments
+  | TF1 Arg Name       | TF2 Arg Name   | Note                             |
+  | ------------------ | -------------  | -------------------------------  |
+  | `learning_rate`    | `learning_rate`| Be careful of setting           |
+  : : : learning_rate tensor value computed from the global step.          :
+  : : : In TF1 this was usually meant to imply a dynamic learning rate and :
+  : : : would recompute in each step. In TF2 (eager + function) it will    :
+  : : : treat it as a scalar value that only gets computed once instead of :
+  : : : a symbolic placeholder to be computed each time.                   :
+  | `decay`            | `rho`          | -                                |
+  | `momentum`         | `momentum`     | -                                |
+  | `epsilon`          | `epsilon`      | Default value is 1e-10 in TF1,   |
+  :                    :                : but 1e-07 in TF2.                :
+  | `use_locking`      | -              | Not applicable in TF2.           |
+
+  #### Before & after usage example
+  Before:
+
+  ```python
+  x = tf.Variable([1,2,3], dtype=tf.float32)
+  grad = tf.constant([0.1, 0.2, 0.3])
+  optimizer = tf.compat.v1.train.RMSPropOptimizer(learning_rate=0.001)
+  optimizer.apply_gradients(zip([grad], [x]))
+  ```
+
+  After:
+
+  ```python
+  x = tf.Variable([1,2,3], dtype=tf.float32)
+  grad = tf.constant([0.1, 0.2, 0.3])
+  optimizer = tf.keras.optimizers.RMSprop(learning_rate=0.001)
+  optimizer.apply_gradients(zip([grad], [x]))
+  ```
+
+  @end_compatibility
+  """
+
+  def __init__(self,
+               learning_rate,
+               decay=0.9,
+               momentum=0.0,
+               epsilon=1e-10,
+               use_locking=False,
+               centered=False,
+               name="RMSProp"):
+    """Construct a new RMSProp optimizer.
+
+    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.
+
+    Args:
+      learning_rate: A Tensor or a floating point value.  The learning rate.
+      decay: Discounting factor for the history/coming gradient
+      momentum: A scalar tensor.
+      epsilon: Small value to avoid zero denominator.
+      use_locking: If True use locks for update operation.
+      centered: 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".
+
+    """
+    super(RMSPropOptimizer, self).__init__(use_locking, name)
+    self._learning_rate = learning_rate
+    self._decay = decay
+    self._momentum = momentum
+    self._epsilon = epsilon
+    self._centered = centered
+
+    # Tensors for learning rate and momentum.  Created in _prepare.
+    self._learning_rate_tensor = None
+    self._decay_tensor = None
+    self._momentum_tensor = None
+    self._epsilon_tensor = None
+
+  def _create_slots(self, var_list):
+    for v in var_list:
+      if v.get_shape().is_fully_defined():
+        init_rms = init_ops.ones_initializer(dtype=v.dtype.base_dtype)
+      else:
+        init_rms = array_ops.ones_like(v)
+      self._get_or_make_slot_with_initializer(v, init_rms, v.get_shape(),
+                                              v.dtype.base_dtype, "rms",
+                                              self._name)
+      if self._centered:
+        self._zeros_slot(v, "mg", self._name)
+      self._zeros_slot(v, "momentum", self._name)
+
+  def _prepare(self):
+    lr = self._call_if_callable(self._learning_rate)
+    decay = self._call_if_callable(self._decay)
+    momentum = self._call_if_callable(self._momentum)
+    epsilon = self._call_if_callable(self._epsilon)
+
+    self._learning_rate_tensor = ops.convert_to_tensor(lr, name="learning_rate")
+    self._decay_tensor = ops.convert_to_tensor(decay, name="decay")
+    self._momentum_tensor = ops.convert_to_tensor(momentum, name="momentum")
+    self._epsilon_tensor = ops.convert_to_tensor(epsilon, name="epsilon")
+
+  def _apply_dense(self, grad, var):
+    rms = self.get_slot(var, "rms")
+    mom = self.get_slot(var, "momentum")
+    if self._centered:
+      mg = self.get_slot(var, "mg")
+      return gen_training_ops.apply_centered_rms_prop(
+          var,
+          mg,
+          rms,
+          mom,
+          math_ops.cast(self._learning_rate_tensor, var.dtype.base_dtype),
+          math_ops.cast(self._decay_tensor, var.dtype.base_dtype),
+          math_ops.cast(self._momentum_tensor, var.dtype.base_dtype),
+          math_ops.cast(self._epsilon_tensor, var.dtype.base_dtype),
+          grad,
+          use_locking=self._use_locking).op
+    else:
+      return gen_training_ops.apply_rms_prop(
+          var,
+          rms,
+          mom,
+          math_ops.cast(self._learning_rate_tensor, var.dtype.base_dtype),
+          math_ops.cast(self._decay_tensor, var.dtype.base_dtype),
+          math_ops.cast(self._momentum_tensor, var.dtype.base_dtype),
+          math_ops.cast(self._epsilon_tensor, var.dtype.base_dtype),
+          grad,
+          use_locking=self._use_locking).op
+
+  def _resource_apply_dense(self, grad, var):
+    rms = self.get_slot(var, "rms")
+    mom = self.get_slot(var, "momentum")
+    if self._centered:
+      mg = self.get_slot(var, "mg")
+      return gen_training_ops.resource_apply_centered_rms_prop(
+          var.handle,
+          mg.handle,
+          rms.handle,
+          mom.handle,
+          math_ops.cast(self._learning_rate_tensor, grad.dtype.base_dtype),
+          math_ops.cast(self._decay_tensor, grad.dtype.base_dtype),
+          math_ops.cast(self._momentum_tensor, grad.dtype.base_dtype),
+          math_ops.cast(self._epsilon_tensor, grad.dtype.base_dtype),
+          grad,
+          use_locking=self._use_locking)
+    else:
+      return gen_training_ops.resource_apply_rms_prop(
+          var.handle,
+          rms.handle,
+          mom.handle,
+          math_ops.cast(self._learning_rate_tensor, grad.dtype.base_dtype),
+          math_ops.cast(self._decay_tensor, grad.dtype.base_dtype),
+          math_ops.cast(self._momentum_tensor, grad.dtype.base_dtype),
+          math_ops.cast(self._epsilon_tensor, grad.dtype.base_dtype),
+          grad,
+          use_locking=self._use_locking)
+
+  def _apply_sparse(self, grad, var):
+    rms = self.get_slot(var, "rms")
+    mom = self.get_slot(var, "momentum")
+    if self._centered:
+      mg = self.get_slot(var, "mg")
+      return gen_training_ops.sparse_apply_centered_rms_prop(
+          var,
+          mg,
+          rms,
+          mom,
+          math_ops.cast(self._learning_rate_tensor, var.dtype.base_dtype),
+          math_ops.cast(self._decay_tensor, var.dtype.base_dtype),
+          math_ops.cast(self._momentum_tensor, var.dtype.base_dtype),
+          math_ops.cast(self._epsilon_tensor, var.dtype.base_dtype),
+          grad.values,
+          grad.indices,
+          use_locking=self._use_locking)
+    else:
+      return gen_training_ops.sparse_apply_rms_prop(
+          var,
+          rms,
+          mom,
+          math_ops.cast(self._learning_rate_tensor, var.dtype.base_dtype),
+          math_ops.cast(self._decay_tensor, var.dtype.base_dtype),
+          math_ops.cast(self._momentum_tensor, var.dtype.base_dtype),
+          math_ops.cast(self._epsilon_tensor, var.dtype.base_dtype),
+          grad.values,
+          grad.indices,
+          use_locking=self._use_locking)
+
+  def _resource_apply_sparse(self, grad, var, indices):
+    rms = self.get_slot(var, "rms")
+    mom = self.get_slot(var, "momentum")
+    if self._centered:
+      mg = self.get_slot(var, "mg")
+      return gen_training_ops.resource_sparse_apply_centered_rms_prop(
+          var.handle,
+          mg.handle,
+          rms.handle,
+          mom.handle,
+          math_ops.cast(self._learning_rate_tensor, grad.dtype),
+          math_ops.cast(self._decay_tensor, grad.dtype),
+          math_ops.cast(self._momentum_tensor, grad.dtype),
+          math_ops.cast(self._epsilon_tensor, grad.dtype),
+          grad,
+          indices,
+          use_locking=self._use_locking)
+    else:
+      return gen_training_ops.resource_sparse_apply_rms_prop(
+          var.handle,
+          rms.handle,
+          mom.handle,
+          math_ops.cast(self._learning_rate_tensor, grad.dtype),
+          math_ops.cast(self._decay_tensor, grad.dtype),
+          math_ops.cast(self._momentum_tensor, grad.dtype),
+          math_ops.cast(self._epsilon_tensor, grad.dtype),
+          grad,
+          indices,
+          use_locking=self._use_locking)

tensorflow/python/training/rmsprop_test.py

@@ -0,0 +1,510 @@
+# 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.
+# ==============================================================================
+"""Tests for rmsprop."""
+
+import copy
+import itertools
+import math
+
+import numpy as np
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import embedding_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import rmsprop
+
+_DATA_TYPES = [dtypes.half, dtypes.float32]
+
+_TEST_PARAM_VALUES = [
+    # learning_rate, decay, momentum, epsilon, centered, use_resource
+    [0.5, 0.9, 0.0, 1e-3, True, False],
+    [0.5, 0.9, 0.0, 1e-3, False, False],
+    [0.5, 0.9, 0.0, 1e-3, True, True],
+    [0.5, 0.9, 0.0, 1e-3, False, True],
+    [0.1, 0.9, 0.0, 1e-3, True, False],
+    [0.5, 0.95, 0.0, 1e-3, False, False],
+    [0.5, 0.95, 0.0, 1e-5, True, False],
+    [0.5, 0.95, 0.9, 1e-5, True, False],
+]
+
+_TESTPARAMS = [
+    [data_type] + values
+    for data_type, values in itertools.product(_DATA_TYPES, _TEST_PARAM_VALUES)
+]
+
+
+class RMSPropOptimizerTest(test.TestCase):
+
+  def _rmsprop_update_numpy(self, var, g, mg, rms, mom, lr, decay, momentum,
+                            epsilon, centered):
+    rms_t = rms * decay + (1 - decay) * g * g
+    denom_t = rms_t + epsilon
+    if centered:
+      mg_t = mg * decay + (1 - decay) * g
+      denom_t -= mg_t * mg_t
+    else:
+      mg_t = mg
+    mom_t = momentum * mom + lr * g / np.sqrt(denom_t, dtype=denom_t.dtype)
+    var_t = var - mom_t
+    return var_t, mg_t, rms_t, mom_t
+
+  def _sparse_rmsprop_update_numpy(self, var, gindexs, gvalues, mg, rms, mom,
+                                   lr, decay, momentum, epsilon, centered):
+    mg_t = copy.deepcopy(mg)
+    rms_t = copy.deepcopy(rms)
+    mom_t = copy.deepcopy(mom)
+    var_t = copy.deepcopy(var)
+    for i in range(len(gindexs)):
+      gindex = gindexs[i]
+      gvalue = gvalues[i]
+      rms_t[gindex] = rms[gindex] * decay + (1 - decay) * gvalue * gvalue
+      denom_t = rms_t[gindex] + epsilon
+      if centered:
+        mg_t[gindex] = mg_t[gindex] * decay + (1 - decay) * gvalue
+        denom_t -= mg_t[gindex] * mg_t[gindex]
+      mom_t[gindex] = momentum * mom[gindex] + lr * gvalue / np.sqrt(denom_t)
+      var_t[gindex] = var[gindex] - mom_t[gindex]
+    return var_t, mg_t, rms_t, mom_t
+
+  @test_util.run_deprecated_v1
+  def testDense(self):
+    # TODO(yori): Use ParameterizedTest when available
+    for (dtype, learning_rate, decay, momentum,
+         epsilon, centered, use_resource) in _TESTPARAMS:
+      with test_util.use_gpu():
+        # Initialize variables for numpy implementation.
+        var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
+        grads0_np = np.array([0.1, 0.2], dtype=dtype.as_numpy_dtype)
+        var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
+        grads1_np = np.array([0.01, 0.2], dtype=dtype.as_numpy_dtype)
+
+        if use_resource:
+          var0 = resource_variable_ops.ResourceVariable(var0_np)
+          var1 = resource_variable_ops.ResourceVariable(var1_np)
+        else:
+          var0 = variables.Variable(var0_np)
+          var1 = variables.Variable(var1_np)
+        grads0 = constant_op.constant(grads0_np)
+        grads1 = constant_op.constant(grads1_np)
+        opt = rmsprop.RMSPropOptimizer(
+            learning_rate=learning_rate,
+            decay=decay,
+            momentum=momentum,
+            epsilon=epsilon,
+            centered=centered)
+
+        update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+        self.evaluate(variables.global_variables_initializer())
+
+        mg0 = opt.get_slot(var0, "mg")
+        self.assertEqual(mg0 is not None, centered)
+        mg1 = opt.get_slot(var1, "mg")
+        self.assertEqual(mg1 is not None, centered)
+        rms0 = opt.get_slot(var0, "rms")
+        self.assertTrue(rms0 is not None)
+        rms1 = opt.get_slot(var1, "rms")
+        self.assertTrue(rms1 is not None)
+        mom0 = opt.get_slot(var0, "momentum")
+        self.assertTrue(mom0 is not None)
+        mom1 = opt.get_slot(var1, "momentum")
+        self.assertTrue(mom1 is not None)
+
+        mg0_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
+        mg1_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
+        rms0_np = np.array([1.0, 1.0], dtype=dtype.as_numpy_dtype)
+        rms1_np = np.array([1.0, 1.0], dtype=dtype.as_numpy_dtype)
+        mom0_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
+        mom1_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
+
+        # Fetch params to validate initial values
+        self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+        self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+
+        # Run 4 steps of RMSProp
+        for _ in range(1, 5):
+          self.evaluate(update)
+
+          var0_np, mg0_np, rms0_np, mom0_np = self._rmsprop_update_numpy(
+              var0_np, grads0_np, mg0_np, rms0_np, mom0_np, learning_rate,
+              decay, momentum, epsilon, centered)
+          var1_np, mg1_np, rms1_np, mom1_np = self._rmsprop_update_numpy(
+              var1_np, grads1_np, mg1_np, rms1_np, mom1_np, learning_rate,
+              decay, momentum, epsilon, centered)
+
+          # Validate updated params
+          if centered:
+            self.assertAllCloseAccordingToType(mg0_np, self.evaluate(mg0))
+            self.assertAllCloseAccordingToType(mg1_np, self.evaluate(mg1))
+          self.assertAllCloseAccordingToType(rms0_np, self.evaluate(rms0))
+          self.assertAllCloseAccordingToType(rms1_np, self.evaluate(rms1))
+          self.assertAllCloseAccordingToType(mom0_np, self.evaluate(mom0))
+          self.assertAllCloseAccordingToType(mom1_np, self.evaluate(mom1))
+          self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0))
+          self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1))
+
+  @test_util.run_deprecated_v1
+  def testMinimizeSparseResourceVariable(self):
+    for dtype in [dtypes.float32, dtypes.float64]:
+      with self.cached_session():
+        var0 = resource_variable_ops.ResourceVariable([[1.0, 2.0]], dtype=dtype)
+        x = constant_op.constant([[4.0], [5.0]], dtype=dtype)
+        pred = math_ops.matmul(embedding_ops.embedding_lookup([var0], [0]), x)
+        loss = pred * pred
+        sgd_op = rmsprop.RMSPropOptimizer(
+            learning_rate=1.0,
+            decay=0.0,
+            momentum=0.0,
+            epsilon=0.0,
+            centered=False).minimize(loss)
+        self.evaluate(variables.global_variables_initializer())
+        # Fetch params to validate initial values
+        self.assertAllCloseAccordingToType([[1.0, 2.0]], self.evaluate(var0))
+        # Run 1 step of sgd
+        self.evaluate(sgd_op)
+        # Validate updated params
+        self.assertAllCloseAccordingToType([[0., 1.]],
+                                           self.evaluate(var0),
+                                           atol=0.01)
+
+  @test_util.run_deprecated_v1
+  def testMinimizeSparseResourceVariableCentered(self):
+    for dtype in [dtypes.float32, dtypes.float64]:
+      with self.cached_session():
+        var0 = resource_variable_ops.ResourceVariable([[1.0, 2.0]], dtype=dtype)
+        x = constant_op.constant([[4.0], [5.0]], dtype=dtype)
+        pred = math_ops.matmul(embedding_ops.embedding_lookup([var0], [0]), x)
+        loss = pred * pred
+        sgd_op = rmsprop.RMSPropOptimizer(
+            learning_rate=1.0,
+            decay=0.0,
+            momentum=0.0,
+            epsilon=1.0,
+            centered=True).minimize(loss)
+        self.evaluate(variables.global_variables_initializer())
+        # Fetch params to validate initial values
+        self.assertAllCloseAccordingToType([[1.0, 2.0]], self.evaluate(var0))
+        # Run 1 step of sgd
+        self.evaluate(sgd_op)
+        # Validate updated params
+        self.assertAllCloseAccordingToType([[-111, -138]],
+                                           self.evaluate(var0),
+                                           atol=0.01)
+
+  @test_util.run_deprecated_v1
+  def testSparse(self):
+    # TODO(yori): Use ParameterizedTest when available
+    for (dtype, learning_rate, decay,
+         momentum, epsilon, centered, _) in _TESTPARAMS:
+      with test_util.use_gpu():
+        # Initialize variables for numpy implementation.
+        var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
+        grads0_np = np.array([0.1], dtype=dtype.as_numpy_dtype)
+        var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
+        grads1_np = np.array([0.01], dtype=dtype.as_numpy_dtype)
+
+        var0 = variables.Variable(var0_np)
+        var1 = variables.Variable(var1_np)
+        grads0_np_indices = np.array([0], dtype=np.int32)
+        grads0 = indexed_slices.IndexedSlices(
+            constant_op.constant(grads0_np),
+            constant_op.constant(grads0_np_indices), constant_op.constant([1]))
+        grads1_np_indices = np.array([1], dtype=np.int32)
+        grads1 = indexed_slices.IndexedSlices(
+            constant_op.constant(grads1_np),
+            constant_op.constant(grads1_np_indices), constant_op.constant([1]))
+        opt = rmsprop.RMSPropOptimizer(
+            learning_rate=learning_rate,
+            decay=decay,
+            momentum=momentum,
+            epsilon=epsilon,
+            centered=centered)
+        update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+        self.evaluate(variables.global_variables_initializer())
+
+        mg0 = opt.get_slot(var0, "mg")
+        self.assertEqual(mg0 is not None, centered)
+        mg1 = opt.get_slot(var1, "mg")
+        self.assertEqual(mg1 is not None, centered)
+        rms0 = opt.get_slot(var0, "rms")
+        self.assertTrue(rms0 is not None)
+        rms1 = opt.get_slot(var1, "rms")
+        self.assertTrue(rms1 is not None)
+        mom0 = opt.get_slot(var0, "momentum")
+        self.assertTrue(mom0 is not None)
+        mom1 = opt.get_slot(var1, "momentum")
+        self.assertTrue(mom1 is not None)
+
+        mg0_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
+        mg1_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
+        rms0_np = np.array([1.0, 1.0], dtype=dtype.as_numpy_dtype)
+        rms1_np = np.array([1.0, 1.0], dtype=dtype.as_numpy_dtype)
+        mom0_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
+        mom1_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
+
+        # Fetch params to validate initial values
+        self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+        self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+
+        # Run 4 steps of RMSProp
+        for _ in range(1, 5):
+          self.evaluate(update)
+
+          var0_np, mg0_np, rms0_np, mom0_np = self._sparse_rmsprop_update_numpy(
+              var0_np, grads0_np_indices, grads0_np, mg0_np, rms0_np, mom0_np,
+              learning_rate, decay, momentum, epsilon, centered)
+          var1_np, mg1_np, rms1_np, mom1_np = self._sparse_rmsprop_update_numpy(
+              var1_np, grads1_np_indices, grads1_np, mg1_np, rms1_np, mom1_np,
+              learning_rate, decay, momentum, epsilon, centered)
+
+          # Validate updated params
+          if centered:
+            self.assertAllCloseAccordingToType(mg0_np, self.evaluate(mg0))
+            self.assertAllCloseAccordingToType(mg1_np, self.evaluate(mg1))
+          self.assertAllCloseAccordingToType(rms0_np, self.evaluate(rms0))
+          self.assertAllCloseAccordingToType(rms1_np, self.evaluate(rms1))
+          self.assertAllCloseAccordingToType(mom0_np, self.evaluate(mom0))
+          self.assertAllCloseAccordingToType(mom1_np, self.evaluate(mom1))
+          self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0))
+          self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1))
+
+  @test_util.run_deprecated_v1
+  def testWithoutMomentum(self):
+    for dtype in [dtypes.half, dtypes.float32]:
+      with test_util.use_gpu():
+        var0 = variables.Variable([1.0, 2.0], dtype=dtype)
+        var1 = variables.Variable([3.0, 4.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.1], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.01], dtype=dtype)
+        opt = rmsprop.RMSPropOptimizer(
+            learning_rate=2.0, decay=0.9, momentum=0.0, epsilon=1.0)
+        update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+        self.evaluate(variables.global_variables_initializer())
+
+        rms0 = opt.get_slot(var0, "rms")
+        self.assertTrue(rms0 is not None)
+        rms1 = opt.get_slot(var1, "rms")
+        self.assertTrue(rms1 is not None)
+        mom0 = opt.get_slot(var0, "momentum")
+        self.assertTrue(mom0 is not None)
+        mom1 = opt.get_slot(var1, "momentum")
+        self.assertTrue(mom1 is not None)
+
+        # 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 rms accumulators where 1. So we should see a normal
+        # update: v -= grad * learning_rate
+        self.evaluate(update)
+        # Check the root mean square accumulators.
+        self.assertAllCloseAccordingToType(
+            np.array([0.901, 0.901]), self.evaluate(rms0))
+        self.assertAllCloseAccordingToType(
+            np.array([0.90001, 0.90001]), self.evaluate(rms1))
+        # Check the parameters.
+        self.assertAllCloseAccordingToType(
+            np.array([
+                1.0 - (0.1 * 2.0 / math.sqrt(0.901 + 1.0)),
+                2.0 - (0.1 * 2.0 / math.sqrt(0.901 + 1.0))
+            ]), self.evaluate(var0))
+        self.assertAllCloseAccordingToType(
+            np.array([
+                3.0 - (0.01 * 2.0 / math.sqrt(0.90001 + 1.0)),
+                4.0 - (0.01 * 2.0 / math.sqrt(0.90001 + 1.0))
+            ]), self.evaluate(var1))
+        # Step 2: the root mean square accumulators contain the previous update.
+        self.evaluate(update)
+        # Check the rms accumulators.
+        self.assertAllCloseAccordingToType(
+            np.array([0.901 * 0.9 + 0.001, 0.901 * 0.9 + 0.001]),
+            self.evaluate(rms0))
+        self.assertAllCloseAccordingToType(
+            np.array([0.90001 * 0.9 + 1e-5, 0.90001 * 0.9 + 1e-5]),
+            self.evaluate(rms1))
+        # Check the parameters.
+        self.assertAllCloseAccordingToType(
+            np.array([
+                1.0 - (0.1 * 2.0 / math.sqrt(0.901 + 1.0)) -
+                (0.1 * 2.0 / math.sqrt(0.901 * 0.9 + 0.001 + 1.0)),
+                2.0 - (0.1 * 2.0 / math.sqrt(0.901 + 1.0)) -
+                (0.1 * 2.0 / math.sqrt(0.901 * 0.9 + 0.001 + 1.0))
+            ]), self.evaluate(var0))
+        self.assertAllCloseAccordingToType(
+            np.array([
+                3.0 - (0.01 * 2.0 / math.sqrt(0.90001 + 1.0)) -
+                (0.01 * 2.0 / math.sqrt(0.90001 * 0.9 + 1e-5 + 1.0)),
+                4.0 - (0.01 * 2.0 / math.sqrt(0.90001 + 1.0)) -
+                (0.01 * 2.0 / math.sqrt(0.90001 * 0.9 + 1e-5 + 1.0))
+            ]), self.evaluate(var1))
+
+  @test_util.run_deprecated_v1
+  def testWithMomentum(self):
+    for dtype in [dtypes.half, dtypes.float32]:
+      with test_util.use_gpu():
+        var0 = variables.Variable([1.0, 2.0], dtype=dtype)
+        var1 = variables.Variable([3.0, 4.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.1], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.01], dtype=dtype)
+
+        opt = rmsprop.RMSPropOptimizer(
+            learning_rate=2.0, decay=0.9, momentum=0.5, epsilon=1e-5)
+        update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+        self.evaluate(variables.global_variables_initializer())
+
+        rms0 = opt.get_slot(var0, "rms")
+        self.assertTrue(rms0 is not None)
+        rms1 = opt.get_slot(var1, "rms")
+        self.assertTrue(rms1 is not None)
+        mom0 = opt.get_slot(var0, "momentum")
+        self.assertTrue(mom0 is not None)
+        mom1 = opt.get_slot(var1, "momentum")
+        self.assertTrue(mom1 is not None)
+
+        # 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: rms = 1, mom = 0. So we should see a normal
+        # update: v -= grad * learning_rate
+        self.evaluate(update)
+        # Check the root mean square accumulators.
+        self.assertAllCloseAccordingToType(
+            np.array([0.901, 0.901]), self.evaluate(rms0))
+        self.assertAllCloseAccordingToType(
+            np.array([0.90001, 0.90001]), self.evaluate(rms1))
+        # Check the momentum accumulators
+        self.assertAllCloseAccordingToType(
+            np.array([(0.1 * 2.0 / math.sqrt(0.901 + 1e-5)),
+                      (0.1 * 2.0 / math.sqrt(0.901 + 1e-5))]),
+            self.evaluate(mom0))
+        self.assertAllCloseAccordingToType(
+            np.array([(0.01 * 2.0 / math.sqrt(0.90001 + 1e-5)),
+                      (0.01 * 2.0 / math.sqrt(0.90001 + 1e-5))]),
+            self.evaluate(mom1))
+
+        # Check that the parameters.
+        self.assertAllCloseAccordingToType(
+            np.array([
+                1.0 - (0.1 * 2.0 / math.sqrt(0.901 + 1e-5)),
+                2.0 - (0.1 * 2.0 / math.sqrt(0.901 + 1e-5))
+            ]), self.evaluate(var0))
+        self.assertAllCloseAccordingToType(
+            np.array([
+                3.0 - (0.01 * 2.0 / math.sqrt(0.90001 + 1e-5)),
+                4.0 - (0.01 * 2.0 / math.sqrt(0.90001 + 1e-5))
+            ]), self.evaluate(var1))
+
+        # Step 2: the root mean square accumulators contain the previous update.
+        self.evaluate(update)
+        # Check the rms accumulators.
+        self.assertAllCloseAccordingToType(
+            np.array([0.901 * 0.9 + 0.001, 0.901 * 0.9 + 0.001]),
+            self.evaluate(rms0))
+        self.assertAllCloseAccordingToType(
+            np.array([0.90001 * 0.9 + 1e-5, 0.90001 * 0.9 + 1e-5]),
+            self.evaluate(rms1))
+        self.assertAllCloseAccordingToType(
+            np.array([
+                0.5 * (0.1 * 2.0 / math.sqrt(0.901 + 1e-5)) +
+                (0.1 * 2.0 / math.sqrt(0.901 * 0.9 + 0.001 + 1e-5)),
+                0.5 * (0.1 * 2.0 / math.sqrt(0.901 + 1e-5)) +
+                (0.1 * 2.0 / math.sqrt(0.901 * 0.9 + 0.001 + 1e-5))
+            ]), self.evaluate(mom0))
+        self.assertAllCloseAccordingToType(
+            np.array([
+                0.5 * (0.01 * 2.0 / math.sqrt(0.90001 + 1e-5)) +
+                (0.01 * 2.0 / math.sqrt(0.90001 * 0.9 + 2e-5)),
+                0.5 * (0.01 * 2.0 / math.sqrt(0.90001 + 1e-5)) +
+                (0.01 * 2.0 / math.sqrt(0.90001 * 0.9 + 2e-5))
+            ]), self.evaluate(mom1))
+
+        # Check the parameters.
+        self.assertAllCloseAccordingToType(
+            np.array([
+                1.0 - (0.1 * 2.0 / math.sqrt(0.901 + 1e-5)) -
+                (0.5 * (0.1 * 2.0 / math.sqrt(0.901 + 1e-5)) +
+                 (0.1 * 2.0 / math.sqrt(0.901 * 0.9 + 0.001 + 1e-5))),
+                2.0 - (0.1 * 2.0 / math.sqrt(0.901 + 1e-5)) -
+                (0.5 * (0.1 * 2.0 / math.sqrt(0.901 + 1e-5)) +
+                 (0.1 * 2.0 / math.sqrt(0.901 * 0.9 + 0.001 + 1e-5)))
+            ]), self.evaluate(var0))
+
+        self.assertAllCloseAccordingToType(
+            np.array([
+                3.0 - (0.01 * 2.0 / math.sqrt(0.90001 + 1e-5)) -
+                (0.5 * (0.01 * 2.0 / math.sqrt(0.90001 + 1e-5)) +
+                 (0.01 * 2.0 / math.sqrt(0.90001 * 0.9 + 2e-5))),
+                4.0 - (0.01 * 2.0 / math.sqrt(0.90001 + 1e-5)) -
+                (0.5 * (0.01 * 2.0 / math.sqrt(0.90001 + 1e-5)) +
+                 (0.01 * 2.0 / math.sqrt(0.90001 * 0.9 + 2e-5)))
+            ]), self.evaluate(var1))
+
+  def testCallableParams(self):
+    with context.eager_mode():
+      for dtype in [dtypes.half, dtypes.float32]:
+        var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
+        var1 = resource_variable_ops.ResourceVariable([3.0, 4.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.1], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.01], dtype=dtype)
+
+        learning_rate = lambda: 2.0
+        decay = lambda: 0.9
+        momentum = lambda: 0.0
+        epsilon = lambda: 1.0
+        opt = rmsprop.RMSPropOptimizer(learning_rate, decay, momentum, epsilon)
+
+        # 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 rms accumulators where 1. So we should see a normal
+        # update: v -= grad * learning_rate
+        opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+        # Check the parameters.
+        self.assertAllCloseAccordingToType(
+            np.array([
+                1.0 - (0.1 * 2.0 / math.sqrt(0.901 + 1.0)),
+                2.0 - (0.1 * 2.0 / math.sqrt(0.901 + 1.0))
+            ]), self.evaluate(var0))
+        self.assertAllCloseAccordingToType(
+            np.array([
+                3.0 - (0.01 * 2.0 / math.sqrt(0.90001 + 1.0)),
+                4.0 - (0.01 * 2.0 / math.sqrt(0.90001 + 1.0))
+            ]), self.evaluate(var1))
+        # Step 2: the root mean square accumulators contain the previous update.
+        opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+        # Check the parameters.
+        self.assertAllCloseAccordingToType(
+            np.array([
+                1.0 - (0.1 * 2.0 / math.sqrt(0.901 + 1.0)) -
+                (0.1 * 2.0 / math.sqrt(0.901 * 0.9 + 0.001 + 1.0)),
+                2.0 - (0.1 * 2.0 / math.sqrt(0.901 + 1.0)) -
+                (0.1 * 2.0 / math.sqrt(0.901 * 0.9 + 0.001 + 1.0))
+            ]), self.evaluate(var0))
+        self.assertAllCloseAccordingToType(
+            np.array([
+                3.0 - (0.01 * 2.0 / math.sqrt(0.90001 + 1.0)) -
+                (0.01 * 2.0 / math.sqrt(0.90001 * 0.9 + 1e-5 + 1.0)),
+                4.0 - (0.01 * 2.0 / math.sqrt(0.90001 + 1.0)) -
+                (0.01 * 2.0 / math.sqrt(0.90001 * 0.9 + 1e-5 + 1.0))
+            ]), self.evaluate(var1))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/saver_large_partitioned_variable_test.py

@@ -0,0 +1,57 @@
+# Copyright 2016 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.
+# =============================================================================
+"""Tests for tensorflow.python.training.saver.py."""
+
+import os
+
+from tensorflow.python.client import session
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import partitioned_variables
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import saver
+
+
+class SaverLargePartitionedVariableTest(test.TestCase):
+
+  # Need to do this in a separate test because of the amount of memory needed
+  # to run this test.
+  def testLargePartitionedVariables(self):
+    save_path = os.path.join(self.get_temp_dir(), "large_variable")
+    var_name = "my_var"
+    # Saving large partition variable.
+    with session.Session("", graph=ops.Graph()) as sess:
+      with ops.device("/cpu:0"):
+        # Create a partitioned variable which is larger than int32 size but
+        # split into smaller sized variables.
+        init = lambda shape, dtype, partition_info: constant_op.constant(
+            True, dtype, shape)
+        partitioned_var = list(variable_scope.get_variable(
+            var_name,
+            shape=[1 << 31],
+            partitioner=partitioned_variables.fixed_size_partitioner(4),
+            initializer=init,
+            dtype=dtypes.bool))
+        self.evaluate(variables.global_variables_initializer())
+        save = saver.Saver(partitioned_var)
+        val = save.save(sess, save_path)
+        self.assertEqual(save_path, val)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/saver_large_variable_test.py

@@ -0,0 +1,56 @@
+# 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.
+# =============================================================================
+"""Tests for tensorflow.python.training.saver.py."""
+
+import os
+
+from tensorflow.core.protobuf import saver_pb2
+from tensorflow.python.client import session
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import saver
+
+
+class SaverLargeVariableTest(test.TestCase):
+
+  # NOTE: This is in a separate file from saver_test.py because the
+  # large allocations do not play well with TSAN, and cause flaky
+  # failures.
+  def testLargeVariable(self):
+    save_path = os.path.join(self.get_temp_dir(), "large_variable")
+    with session.Session("", graph=ops.Graph()) as sess:
+      # Declare a variable that is exactly 2GB. This should fail,
+      # because a serialized checkpoint includes other header
+      # metadata.
+      with ops.device("/cpu:0"):
+        var = variables.Variable(
+            constant_op.constant(
+                False, shape=[2, 1024, 1024, 1024], dtype=dtypes.bool))
+      save = saver.Saver(
+          {
+              var.op.name: var
+          }, write_version=saver_pb2.SaverDef.V1)
+      var.initializer.run()
+      with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
+                                  "Tensor slice is too large to serialize"):
+        save.save(sess, save_path)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/saver_test_utils.py

@@ -0,0 +1,87 @@
+# 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.
+# =============================================================================
+"""Utility classes for testing checkpointing."""
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops as ops_lib
+from tensorflow.python.ops import gen_lookup_ops
+from tensorflow.python.training import saver as saver_module
+
+
+class CheckpointedOp:
+  """Op with a custom checkpointing implementation.
+
+  Defined as part of the test because the MutableHashTable Python code is
+  currently in contrib.
+  """
+
+  # pylint: disable=protected-access
+  def __init__(self, name, table_ref=None):
+    if table_ref is None:
+      self.table_ref = gen_lookup_ops.mutable_hash_table_v2(
+          key_dtype=dtypes.string, value_dtype=dtypes.float32, name=name)
+    else:
+      self.table_ref = table_ref
+    self._name = name
+    if not context.executing_eagerly():
+      self._saveable = CheckpointedOp.CustomSaveable(self, name)
+      ops_lib.add_to_collection(ops_lib.GraphKeys.SAVEABLE_OBJECTS,
+                                self._saveable)
+
+  @property
+  def name(self):
+    return self._name
+
+  @property
+  def saveable(self):
+    if context.executing_eagerly():
+      return CheckpointedOp.CustomSaveable(self, self.name)
+    else:
+      return self._saveable
+
+  def insert(self, keys, values):
+    return gen_lookup_ops.lookup_table_insert_v2(self.table_ref, keys, values)
+
+  def lookup(self, keys, default):
+    return gen_lookup_ops.lookup_table_find_v2(self.table_ref, keys, default)
+
+  def keys(self):
+    return self._export()[0]
+
+  def values(self):
+    return self._export()[1]
+
+  def _export(self):
+    return gen_lookup_ops.lookup_table_export_v2(self.table_ref, dtypes.string,
+                                                 dtypes.float32)
+
+  class CustomSaveable(saver_module.BaseSaverBuilder.SaveableObject):
+    """A custom saveable for CheckpointedOp."""
+
+    def __init__(self, table, name):
+      tensors = table._export()
+      specs = [
+          saver_module.BaseSaverBuilder.SaveSpec(tensors[0], "",
+                                                 name + "-keys"),
+          saver_module.BaseSaverBuilder.SaveSpec(tensors[1], "",
+                                                 name + "-values")
+      ]
+      super(CheckpointedOp.CustomSaveable, self).__init__(table, specs, name)
+
+    def restore(self, restore_tensors, shapes):
+      return gen_lookup_ops.lookup_table_import_v2(
+          self.op.table_ref, restore_tensors[0], restore_tensors[1])
+  # pylint: enable=protected-access

tensorflow/python/training/saving/BUILD

@@ -0,0 +1,106 @@
+# Description:
+#   Low-level utilities for reading and writing checkpoints.
+
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//tensorflow:internal",
+    ],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "checkpoint_options",
+    srcs = ["checkpoint_options.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/checkpoint:checkpoint_options",
+        "//tensorflow/python/util:deprecation",
+    ],
+)
+
+py_library(
+    name = "functional_saver",
+    srcs = ["functional_saver.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/checkpoint:functional_saver",
+        "//tensorflow/python/util:deprecation",
+    ],
+)
+
+py_library(
+    name = "saveable_object",
+    srcs = ["saveable_object.py"],
+    strict_deps = True,
+)
+
+py_library(
+    name = "saveable_object_util",
+    srcs = ["saveable_object_util.py"],
+    strict_deps = True,
+    deps = [
+        ":saveable_object",
+        "//tensorflow/python/checkpoint:saveable_compat",
+        "//tensorflow/python/client:session",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:device",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:control_flow_ops_gen",
+        "//tensorflow/python/ops:ref_variable",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/ops:state_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/trackable:base",
+        "//tensorflow/python/trackable:base_delegate",
+        "//tensorflow/python/trackable:python_state",
+        "//tensorflow/python/trackable:trackable_utils",
+        "//tensorflow/python/types:core",
+        "//tensorflow/python/util:compat",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:object_identity",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "trace_saveable_util",
+    srcs = ["trace_saveable_util.py"],
+    strict_deps = True,
+    deps = [
+        ":saveable_object",
+        ":saveable_object_util",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/framework:type_spec",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/util:nest",
+    ],
+)
+
+tf_py_strict_test(
+    name = "saveable_object_util_test",
+    srcs = ["saveable_object_util_test.py"],
+    deps = [
+        ":saveable_object",
+        ":saveable_object_util",
+        "//tensorflow/python/checkpoint",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:resource_variable_ops_gen",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/trackable:base",
+        "//tensorflow/python/trackable:resource",
+    ],
+)

tensorflow/python/training/saving/checkpoint_options.py

@@ -0,0 +1,23 @@
+# 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.
+# ==============================================================================
+"""Options for saving Checkpoints."""
+
+# TODO(kathywu): Delete this file after all imports have been moved to the path
+# below.
+from tensorflow.python.checkpoint import checkpoint_options
+from tensorflow.python.util import deprecation
+
+__getattr__ = deprecation.deprecate_moved_module(
+    __name__, checkpoint_options, "2.11")

tensorflow/python/training/saving/functional_saver.py

@@ -0,0 +1,23 @@
+# 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.
+# ==============================================================================
+"""Saves and restore variables inside traced @tf.functions."""
+
+# TODO(kathywu): Delete this file after all imports have been moved to the path
+# below.
+from tensorflow.python.checkpoint import functional_saver
+from tensorflow.python.util import deprecation
+
+__getattr__ = deprecation.deprecate_moved_module(
+    __name__, functional_saver, "2.11")

tensorflow/python/training/saving/saveable_object.py

@@ -0,0 +1,94 @@
+# 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.
+# ==============================================================================
+"""Types for specifying saving and loading behavior."""
+
+
+class SaveSpec:
+  """Class used to describe tensor slices that need to be saved."""
+
+  def __init__(self, tensor, slice_spec, name, dtype=None, device=None):
+    """Creates a `SaveSpec` object.
+
+    Args:
+      tensor: the tensor to save or callable that produces a tensor to save.
+        If the value is `None`, the `SaveSpec` is ignored.
+      slice_spec: the slice to be saved. See `Variable.SaveSliceInfo`.
+      name: the name to save the tensor under.
+      dtype: The data type of the Tensor. Required if `tensor` is callable.
+        Used for error checking in the restore op.
+      device: The device generating and consuming this tensor. Required if
+        `tensor` is callable. Used to group objects to save by device.
+    """
+    self._tensor = tensor
+    self.slice_spec = slice_spec
+    self.name = name
+    if callable(self._tensor):
+      if dtype is None or device is None:
+        raise AssertionError(
+            "When passing a callable `tensor` to a SaveSpec, an explicit "
+            "dtype and device must be provided.")
+      self.dtype = dtype
+      self.device = device
+    else:
+      self.dtype = tensor.dtype
+      if device is not None:
+        self.device = device
+      else:
+        self.device = tensor.device
+
+  @property
+  def tensor(self):
+    return self._tensor() if callable(self._tensor) else self._tensor
+
+
+class SaveableObject:
+  """Base class for saving and restoring saveable objects."""
+
+  def __init__(self, op, specs, name):
+    """Creates a `SaveableObject` object.
+
+    Args:
+      op: the "producer" object that this class wraps; it produces a list of
+        tensors to save.  E.g., a "Variable" object saving its backing tensor.
+      specs: a list of SaveSpec, each element of which describes one tensor to
+        save under this object. All Tensors must be on the same device.
+      name: the name to save the object under.
+    """
+    self.op = op
+    self.specs = specs
+    self.name = name
+
+  @property
+  def device(self):
+    """The device for SaveSpec Tensors."""
+    return self.specs[0].device
+
+  def restore(self, restored_tensors, restored_shapes):
+    """Restores this object from 'restored_tensors'.
+
+    Args:
+      restored_tensors: the tensors that were loaded from a checkpoint
+      restored_shapes: the shapes this object should conform to after
+        restore, or None.
+
+    Returns:
+      An operation that restores the state of the object.
+
+    Raises:
+      ValueError: If the object cannot be restored using the provided
+        parameters.
+    """
+    # pylint: disable=unused-argument
+    raise ValueError("Calling an abstract method.")

tensorflow/python/training/saving/saveable_object_util.py

@@ -0,0 +1,839 @@
+# 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.
+# ==============================================================================
+"""Utilities for working with and creating SaveableObjects."""
+import functools
+
+from tensorflow.python.checkpoint import saveable_compat
+from tensorflow.python.client import session
+from tensorflow.python.eager import context
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import device as pydev
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor as tensor_lib
+from tensorflow.python.framework import tensor_util
+
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_control_flow_ops
+from tensorflow.python.ops import ref_variable
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import state_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.trackable import base as trackable
+from tensorflow.python.trackable import base_delegate
+from tensorflow.python.trackable import python_state
+from tensorflow.python.trackable import trackable_utils
+from tensorflow.python.training.saving import saveable_object
+from tensorflow.python.types import core
+from tensorflow.python.util import compat
+from tensorflow.python.util import nest
+from tensorflow.python.util import object_identity
+from tensorflow.python.util.tf_export import tf_export
+
+# Op names which identify variable reads which should be saved.
+_VARIABLE_OPS = set(["Variable",
+                     "VariableV2",
+                     "AutoReloadVariable",
+                     "VarHandleOp",
+                     "ReadVariableOp"])
+
+_REF_VARIABLE_OPS = frozenset(["Variable", "VariableV2", "AutoReloadVariable"])
+
+
+def set_cpu0(device_string):
+  """Creates a new device string based on `device_string` but using /CPU:0.
+
+  If the device is already on /CPU:0 or it is a custom device, this is a no-op.
+
+  Args:
+    device_string: A device string.
+
+  Returns:
+    A device string.
+  """
+  if context.is_custom_device(device_string):
+    return device_string
+  parsed_device = pydev.DeviceSpec.from_string(device_string)
+  parsed_device = parsed_device.replace(device_type="CPU", device_index=0)
+  return parsed_device.to_string()
+
+
+class ReferenceVariableSaveable(saveable_object.SaveableObject):
+  """SaveableObject implementation that handles reference variables."""
+
+  def __init__(self, var, slice_spec, name):
+    spec = saveable_object.SaveSpec(var, slice_spec, name, dtype=var.dtype)
+    super(ReferenceVariableSaveable, self).__init__(var, [spec], name)
+
+  def restore(self, restored_tensors, restored_shapes):
+    restored_tensor = restored_tensors[0]
+    if restored_shapes is not None:
+      restored_tensor = array_ops.reshape(restored_tensor, restored_shapes[0])
+    return state_ops.assign(
+        self.op,
+        restored_tensor,
+        validate_shape=restored_shapes is None and
+        self.op.get_shape().is_fully_defined())
+
+
+class ResourceVariableSaveable(saveable_object.SaveableObject):
+  """SaveableObject implementation that handles ResourceVariables."""
+
+  def __init__(self, var, slice_spec, name):
+    self._var_device = var.device
+    self._var_shape = var.shape
+    if isinstance(var, tensor_lib.Tensor):
+      self.handle_op = var.op.inputs[0]
+      tensor = var
+    elif resource_variable_ops.is_resource_variable(var):
+
+      def _read_variable_closure(v):
+        def f():
+          with ops.device(v.device):
+            if context.executing_eagerly() and not v.is_initialized():
+              # A SaveSpec tensor value of `None` indicates that the variable is
+              # uninitialized.
+              return None
+            # Read the variable without making a copy to limit memory usage.
+            x = v.read_value_no_copy()
+            # To allow variables placed on non-CPU devices to be checkpointed,
+            # we copy them to CPU on the same machine first.
+            with ops.device("/device:CPU:0"):
+              return array_ops.identity(x)
+
+        return f
+
+      self.handle_op = var.handle
+      tensor = _read_variable_closure(var)
+    else:
+      raise ValueError(
+          "Saveable is neither a resource variable nor a read operation."
+          f" Got: {repr(var)}")
+    spec = saveable_object.SaveSpec(tensor, slice_spec, name,
+                                    dtype=var.dtype, device=var.device)
+    super(ResourceVariableSaveable, self).__init__(var, [spec], name)
+
+  def restore(self, restored_tensors, restored_shapes):
+    """Restores tensors. Raises ValueError if incompatible shape found."""
+    restored_tensor = restored_tensors[0]
+    if restored_shapes is not None:
+      restored_tensor = array_ops.reshape(restored_tensor, restored_shapes[0])
+    # Copy the restored tensor to the variable's device.
+    with ops.device(self._var_device):
+      restored_tensor = array_ops.identity(restored_tensor)
+      try:
+        assigned_variable = resource_variable_ops.shape_safe_assign_variable_handle(
+            self.handle_op, self._var_shape, restored_tensor)
+      except ValueError as e:
+        raise ValueError(
+            f"Received incompatible tensor with shape {restored_tensor.shape} "
+            f"when attempting to restore variable with shape {self._var_shape} "
+            f"and name {self.name}.") from e
+      return assigned_variable
+
+
+def _tensor_comes_from_variable(v):
+  return isinstance(v, tensor_lib.Tensor) and v.op.type in _VARIABLE_OPS
+
+
+def saveable_objects_for_op(op, name):
+  """Create `SaveableObject`s from an operation.
+
+  Args:
+    op: A variable, operation, or SaveableObject to coerce into a
+      SaveableObject.
+    name: A string name for the SaveableObject.
+
+  Yields:
+    `SaveableObject`s which together save/restore `op`.
+
+  Raises:
+    TypeError: If `name` is not a string.
+    ValueError: For operations with no known conversion to SaveableObject.
+  """
+  if not isinstance(name, str):
+    raise TypeError(
+        "names_to_saveables must be a dict mapping string names to "
+        f"trackable operations. Name is not a string: {name}")
+  if isinstance(op, saveable_object.SaveableObject):
+    yield op
+  elif isinstance(op, (list, tuple, variables.PartitionedVariable)):
+    if isinstance(op, variables.PartitionedVariable):
+      op = list(op)
+    # A set of slices.
+    slice_name = None
+    # pylint: disable=protected-access
+    for variable in op:
+      if isinstance(variable, saveable_object.SaveableObject):
+        yield variable
+        continue
+      if not isinstance(variable, variables.Variable):
+        raise ValueError(f"Slices must all be Variables: {variable}")
+      if not variable._save_slice_info:
+        raise ValueError(f"Slices must all be slices: {variable}")
+      if slice_name is None:
+        slice_name = variable._save_slice_info.full_name
+      elif slice_name != variable._save_slice_info.full_name:
+        raise ValueError(
+            f"Slices must all be from the same tensor: {slice_name} != "
+            f"{variable._save_slice_info.full_name}")
+      if variable.op.type in _REF_VARIABLE_OPS:
+        yield ReferenceVariableSaveable(
+            variable, variable._save_slice_info.spec, name)
+      else:
+        yield ResourceVariableSaveable(variable, variable._save_slice_info.spec,
+                                       name)
+    # pylint: enable=protected-access
+  elif isinstance(op, trackable.Trackable) and not isinstance(
+      op, variables.Variable):
+    # pylint: disable=protected-access
+    for attr, factory in saveable_objects_from_trackable(
+        op, tf1_saver=True).items():
+      if attr == trackable.VARIABLE_VALUE_KEY:
+        # Keep original name for classes masquerading as variables and
+        # Trackables that define _serialize_to_tensors.
+        full_name = name
+      elif attr == trackable_utils.SERIALIZE_TO_TENSORS_NAME:
+        full_name = name
+      else:
+        full_name = name + "_" + attr
+      op = (factory(full_name) if callable(factory) else factory)
+      for op in saveable_objects_for_op(op, op.name):
+        yield op
+    # pylint: enable=protected-access
+  else:
+    # A variable or tensor.
+    if isinstance(op, resource_variable_ops.BaseResourceVariable):
+      if op._in_graph_mode:  # pylint: disable=protected-access
+        variable = op._graph_element  # pylint: disable=protected-access
+      else:
+        variable = op
+      yield ResourceVariableSaveable(variable, "", name)
+    else:
+      if context.executing_eagerly():
+        raise ValueError("Can only save/restore ResourceVariables when "
+                         f"executing eagerly, got type: {type(op)}.")
+
+      variable = ops.convert_to_tensor(op, as_ref=True)
+      if not _tensor_comes_from_variable(variable):
+        raise TypeError(
+            "names_to_saveables must be a dict mapping string "
+            f"names to Tensors/Variables. Not a variable: {variable}")
+      if variable.op.type in _REF_VARIABLE_OPS:
+        yield ReferenceVariableSaveable(variable, "", name)
+      else:
+        yield ResourceVariableSaveable(variable, "", name)
+
+
+def op_list_to_dict(op_list, convert_variable_to_tensor=True):
+  """Create a dictionary of names to operation lists.
+
+  This method is only used when the variable name matters (e.g. when saving
+  or restoring from a TF1 name-based checkpoint). In TF2, this can be called
+  from `tf.train.Checkpoint.restore` when loading from a name-based checkpoint.
+
+  Args:
+    op_list: A (nested) list, tuple, or set of Variables or SaveableObjects.
+    convert_variable_to_tensor: Whether or not to convert single Variables
+      with no slice info into Tensors.
+
+  Returns:
+    A dictionary of names to the operations that must be saved under
+    that name.  Variables with save_slice_info are grouped together under the
+    same key in no particular order.
+
+  Raises:
+    TypeError: If the type of op_list or its elements is not supported.
+    ValueError: If at least two saveables share the same name.
+  """
+  if not isinstance(op_list, (list, tuple, set)):
+    raise TypeError("Variables to save should be passed in a dict or a "
+                    f"list. Got {op_list}")
+  # List casting is necessary to support sets.
+  op_list = nest.flatten(list(op_list))
+  # When ResourceVariables are converted to Tensors, read ops are added to the
+  # graph. Sorting the op_list ensures that the resulting graph is always
+  # constructed in a deterministic way:
+  op_list = sorted(op_list, key=lambda x: x.name)
+  names_to_saveables = {}
+  # pylint: disable=protected-access
+  for var in op_list:
+    resource_or_ref_variable = (
+        isinstance(var, resource_variable_ops.BaseResourceVariable) or
+        isinstance(var, ref_variable.RefVariable))
+
+    if isinstance(var, saveable_object.SaveableObject):
+      names_to_saveables[var.name] = var
+    elif isinstance(var, variables.PartitionedVariable):
+      if var.name in names_to_saveables:
+        raise ValueError(
+            f"At least two variables have the same name: {var.name}")
+      names_to_saveables[var.name] = var
+    elif isinstance(var, variables.Variable) and var._save_slice_info:
+      name = var._save_slice_info.full_name
+      if name in names_to_saveables:
+        if not isinstance(names_to_saveables[name], list):
+          raise ValueError("Mixing slices and non-slices with the same name: "
+                           f"{name}")
+        names_to_saveables[name].append(var)
+      else:
+        names_to_saveables[name] = [var]
+    elif isinstance(var, trackable.Trackable) and not resource_or_ref_variable:
+      trackable_saveables = [
+          (factory() if callable(factory) else factory)
+          for factory in (
+              saveable_objects_from_trackable(var, tf1_saver=True).values())]
+      names_to_saveables.update(
+          op_list_to_dict(trackable_saveables))
+    else:
+      # Variables (reference and resource) have an _in_graph_mode property
+      # indicating whether they were created in a graph building context. We
+      # also get Tensors when graph building, which do not have this property.
+      if not getattr(var, "_in_graph_mode", True):
+        if not isinstance(var, resource_variable_ops.BaseResourceVariable):
+          raise ValueError(
+              "Can only save/restore ResourceVariables when eager execution "
+              f"is enabled. Got type: {type(var)}.")
+        set_var = names_to_saveables.setdefault(var._shared_name, var)
+        if set_var is not var:
+          raise ValueError(
+              "Two different ResourceVariable objects with the same "
+              f"shared_name '{var._shared_name}' were passed to the Saver. This"
+              " likely means that they were created in different Graphs or "
+              "isolated contexts, and may not be checkpointed together.")
+      else:
+        if convert_variable_to_tensor:
+          if isinstance(var, resource_variable_ops.BaseResourceVariable):
+            var = var._graph_element  # pylint: disable=protected-access
+          else:
+            var = ops.convert_to_tensor(var, as_ref=True)
+          if not _tensor_comes_from_variable(var):
+            raise TypeError(f"Variable to save is not a Variable: {var}")
+        if var.op.type == "ReadVariableOp":
+          name = var.op.inputs[0].op.name
+        else:
+          name = var.op.name
+        if name in names_to_saveables:
+          raise ValueError(f"At least two variables have the same name: {name}")
+        names_to_saveables[name] = var
+
+    # pylint: enable=protected-access
+  return names_to_saveables
+
+
+def _add_saveable(saveables, seen_ops, saveable):
+  """Adds the saveable to the saveables list.
+
+  Args:
+    saveables: List to append the SaveableObject to.
+    seen_ops: Set of the ops of the saveables already processed.  Used to
+      check that each saveable is only saved once.
+    saveable: The saveable.
+
+  Raises:
+    ValueError: If the saveable has already been processed.
+  """
+  if saveable.op is not None and saveable.op in seen_ops:
+    raise ValueError("The same saveable will be restored with two names: "
+                     f"{saveable.name}")
+  saveables.append(saveable)
+  seen_ops.add(saveable.op)
+
+
+def validate_and_slice_inputs(names_to_saveables):
+  """Returns the variables and names that will be used for a Saver.
+
+  Args:
+    names_to_saveables: A dict (k, v) where k is the name of an operation and
+       v is an operation to save or a BaseSaverBuilder.Saver.
+
+  Returns:
+    A list of SaveableObjects.
+
+  Raises:
+    TypeError: If any of the keys are not strings or any of the
+      values are not one of Tensor or Variable or a trackable operation.
+    ValueError: If the same operation is given in more than one value
+      (this also applies to slices of SlicedVariables).
+  """
+  saveables = []
+  seen_ops = object_identity.ObjectIdentitySet()
+  for name, op in sorted(names_to_saveables.items(),
+                         # Avoid comparing ops, sort only by name.
+                         key=lambda x: x[0]):
+    for converted_saveable_object in saveable_objects_for_op(op, name):
+      _add_saveable(saveables, seen_ops, converted_saveable_object)
+  return saveables
+
+
+def validate_saveables_for_saved_model(saveables, obj):
+  """Makes sure SaveableObjects are compatible with SavedModel."""
+  if isinstance(obj, python_state.PythonState):
+    logging.warn(
+        f"Note that object {obj} stores python values into the checkpoint. "
+        "These values will not be restored when loading the SavedModel "
+        "into python.")
+    return []
+  if any(isinstance(saveable, trackable.NoRestoreSaveable)
+         for saveable in saveables):
+    return []
+  return saveables
+
+
+class RestoredSaveableObject(saveable_object.SaveableObject):
+  """SaveableObject restored from SavedModel using the traced save/restore."""
+
+  def __init__(self, names_and_slices, save_function, restore_function, name):
+    self.save_function = save_function
+    self.restore_function = restore_function
+
+    if tensor_util.is_tf_type(name):
+      name_tensor = name
+    else:
+      with ops.init_scope():
+        name_tensor = constant_op.constant(name)
+    tensors = save_function(name_tensor)
+    specs = []
+    for (str_name, str_slice), tensor_info in zip(names_and_slices, tensors):
+      specs.append(saveable_object.SaveSpec(tensor_info["tensor"], str_slice,
+                                            name + str_name))
+    super(RestoredSaveableObject, self).__init__(None, specs, name)
+
+  def restore(self, restored_tensors, restored_shapes):
+    del restored_shapes  # unused
+    return self.restore_function(
+        *[restored_tensors[i] for i in range(len(self.specs))])
+
+
+def recreate_saveable_objects(saveable_fn_by_name, temp_session):
+  """Returns a dict of SaveableObject factories generated from loaded fns."""
+
+  names_and_slices = []
+
+  with ops.init_scope():
+
+    for save_fn, _ in saveable_fn_by_name.values():
+      for tensor_info in save_fn(""):
+        name = tensor_info["name"]
+        slice_spec = tensor_info["slice_spec"]
+        if not context.executing_eagerly():
+          sess = ops.get_default_session()
+          if sess is None:
+            if temp_session[0] is not None:
+              sess = temp_session[0]
+            else:
+              sess = temp_session[0] = session.Session()
+          name, slice_spec = sess.run([name, slice_spec])
+        names_and_slices.append((
+            _convert_to_string(name),
+            _convert_to_string(slice_spec)))
+
+  saveable_factories = {}
+  for name, (save_fn, restore_fn) in saveable_fn_by_name.items():
+    saveable_factories[name] = functools.partial(
+        RestoredSaveableObject,
+        names_and_slices=names_and_slices,
+        save_function=save_fn,
+        restore_function=restore_fn)
+  return saveable_factories
+
+
+def create_saveable_object(name, key, factory, call_with_mapped_captures):
+  """Creates a SaveableObject while potentially in a different graph.
+
+  When creating the frozen saver for SavedModel, the save and restore ops are
+  placed in a separate graph. Since RestoredSaveableObject uses tf.functions to
+  save and restore, the function captures must be mapped to the new graph.
+
+  Args:
+    name: Name of SaveableObject factory.
+    key: Checkpoint key of this SaveableObject.
+    factory: Factory method for creating the SaveableObject.
+    call_with_mapped_captures: Helper that calls a tf.function while remapping
+      the captures.
+
+  Returns:
+    a SaveableObject.
+  """
+  if call_with_mapped_captures is None:
+    return factory(name=key)
+  if name == trackable_utils.SERIALIZE_TO_TENSORS_NAME:
+    return factory(name=key,
+                   call_with_mapped_captures=call_with_mapped_captures)
+  elif is_factory_for_restored_saveable_object(factory):
+    concrete_save_fn = factory.keywords["save_function"]
+
+    def save_fn(name):
+      return call_with_mapped_captures(concrete_save_fn, [name])
+
+    concrete_restore_fn = factory.keywords["restore_function"]
+
+    def restore_fn(*restored_tensors):
+      return call_with_mapped_captures(concrete_restore_fn, restored_tensors)
+
+    return factory(save_function=save_fn, restore_function=restore_fn,
+                   name=key)
+  else:
+    return factory(name=key)
+
+
+def is_factory_for_restored_saveable_object(factory):
+  return (isinstance(factory, functools.partial) and
+          factory.func is RestoredSaveableObject)
+
+
+@tf_export("__internal__.tracking.saveable_objects_from_trackable", v1=[])
+def saveable_objects_from_trackable(obj, tf1_saver=False):
+  """Returns SaveableObject factory dict from a Trackable.
+
+  Args:
+    obj: A `Trackable`
+    tf1_saver: Boolean, whether this is being called from a TF1 Saver (
+        `tf.compat.v1.train.Saver`). When this is True, the SaveableObject will
+        be generated from `obj`'s legacy `_gather_saveables_for_checkpoint` fn.
+        When saving with TF2, `Trackable._serialize_from_tensors` is preferred.
+
+  Returns:
+    A dict mapping attribute names to SaveableObject factories (callables that
+    produce a SaveableObject).
+  """
+  if isinstance(obj, python_state.PythonState):
+    return {
+        python_state.PYTHON_STATE:
+            functools.partial(
+                _PythonStringStateSaveable,
+                state_callback=obj.serialize,
+                restore_callback=obj.deserialize)
+    }
+
+  if tf1_saver:
+    saveable_factories = obj._gather_saveables_for_checkpoint()  # pylint: disable=protected-access
+    if saveable_factories:
+      return saveable_factories
+
+  if trackable_has_serialize_to_tensor(obj):
+
+    def create_saveable(name="", call_with_mapped_captures=None):
+      save_fn = obj._serialize_to_tensors  # pylint: disable=protected-access
+      if (call_with_mapped_captures and
+          isinstance(save_fn, core.ConcreteFunction)):
+        tensor_dict = call_with_mapped_captures(save_fn, [])
+      else:
+        tensor_dict = save_fn()
+
+      specs = []
+      local_names = []
+      for tensor_name, maybe_tensor in tensor_dict.items():
+        local_names.append(tensor_name)
+
+        if not isinstance(maybe_tensor, dict):
+          maybe_tensor = {"": maybe_tensor}
+
+        spec_name = name + trackable_utils.escape_local_name(tensor_name)
+        # Create separate specs for each slice spec.
+        for slice_spec, tensor in maybe_tensor.items():
+          if isinstance(tensor, saveable_object.SaveSpec):
+            spec = tensor
+            spec.name = spec_name
+            spec.slice_spec = slice_spec
+          else:
+            spec = saveable_object.SaveSpec(tensor, slice_spec, spec_name)
+          specs.append(spec)
+
+      return TrackableSaveable(
+          obj=obj,
+          specs=specs,
+          name=name,
+          local_names=local_names,
+          prefix=saveable_compat.get_saveable_name(obj) or "",
+          call_with_mapped_captures=call_with_mapped_captures)
+
+    return {trackable_utils.SERIALIZE_TO_TENSORS_NAME: create_saveable}
+  else:
+    return obj._gather_saveables_for_checkpoint()  # pylint: disable=protected-access
+
+
+class TrackableSaveable(saveable_object.SaveableObject):
+  """A SaveableObject that defines `Trackable` checkpointing steps."""
+
+  def __init__(self, obj, specs, name, local_names, prefix,
+               call_with_mapped_captures=None):
+    self._prefix = prefix
+    self._local_names = local_names
+    self._trackable = obj
+    self._call_with_mapped_captures = call_with_mapped_captures
+    super(TrackableSaveable, self).__init__(obj, specs, name)
+
+  def restore(self, restored_tensors, restored_shapes):
+    del restored_shapes  # Unused.
+    restored_tensor_dict = {}
+    for n, local_name in enumerate(self._local_names):
+      restored_tensor_dict[local_name] = restored_tensors[n]
+
+    restore_fn = self._trackable._restore_from_tensors  # pylint: disable=protected-access
+
+    # When restoring a RefVariable, call the restore function directly.
+    # pylint: disable=protected-access
+    if not ops.executing_eagerly_outside_functions() and any([
+        spec._tensor.op.type in _REF_VARIABLE_OPS
+        for spec in self.specs
+        if isinstance(spec._tensor, tensor_lib.Tensor)]):
+      return restore_fn(restored_tensor_dict)
+    # pylint: enable=protected-access
+
+    if (self._call_with_mapped_captures and
+        isinstance(restore_fn, core.ConcreteFunction)):
+      ret = self._call_with_mapped_captures(restore_fn, [restored_tensor_dict])
+    else:
+      ret = restore_fn(restored_tensor_dict)
+    if ret is not None:
+      return ret
+    return gen_control_flow_ops.no_op()
+
+  def get_proto_names_and_checkpoint_keys(self):
+    return [(self._prefix + local_name, spec.name)
+            for local_name, spec in zip(self._local_names, self.specs)]
+
+
+class _PythonStringStateSaveable(saveable_object.SaveableObject):
+  """Saves Python state in a checkpoint."""
+
+  def __init__(self, name, state_callback, restore_callback):
+    """Configure saving.
+
+    Args:
+      name: The checkpoint key to write to.
+      state_callback: A function taking no arguments which returns a string.
+        This function is run every time a checkpoint is written.
+      restore_callback: A function taking a Python string, used to restore
+        state.
+    """
+
+    def _state_callback_wrapper():
+      with ops.init_scope():
+        return state_callback()
+
+    self._state_callback = _state_callback_wrapper
+    self._restore_callback = restore_callback
+    with ops.device("/cpu:0"):
+      self._save_string = constant_op.constant("", dtype=dtypes.string)
+    spec = saveable_object.SaveSpec(
+        self._save_string, "", name, dtype=dtypes.string)
+    super(_PythonStringStateSaveable, self).__init__(self._save_string, [spec],
+                                                     name)
+
+  def feed_dict_additions(self):
+    """When running a graph, indicates fresh state to feed."""
+    return {self._save_string: self._state_callback()}
+
+  def freeze(self):
+    """Create a frozen `SaveableObject` which saves the current state."""
+
+    def _constant_state():
+      return constant_op.constant(self._state_callback(), dtype=dtypes.string)
+
+    return trackable.NoRestoreSaveable(
+        tensor=_constant_state,
+        dtype=dtypes.string,
+        name=self.name,
+        device="cpu:0")
+
+
+def trackable_has_serialize_to_tensor(obj):
+  """Returns whether obj's class has `_serialize_to_tensors` defined."""
+  if obj is base_delegate.DelegatingTrackableMixin:
+    # DelegatingTrackableMixin always delegates "_serialize_to_tensors"
+    # to its inner `trackable`, so we check whether the inner trackable
+    # has `_serialize_to_tensor`.
+    return trackable_has_serialize_to_tensor(obj._trackable)  # pylint: disable=protected-access
+
+  try:
+    if "_serialize_to_tensors" in obj.__dict__:
+      # In some cases (e.g. restored objects), the object may have
+      # `_serialize_to_tensors` even if the class does not.
+      return True
+  except (AttributeError, TypeError):
+    # Data structure proxy wrappers don't have __dict__.
+    pass
+
+  # Use MRO so that if a parent class has `_serialize_to_tensors`, but the
+  # object class has not yet been migrated, we'll continue to use the obj
+  # class's `_gather_saveables_for_checkpoint` method.
+  for t in type(obj).mro():
+    if t is base_delegate.DelegatingTrackableMixin:
+      # DelegatingTrackableMixin always delegates "_serialize_to_tensors"
+      # to its inner `trackable`, so we check whether the inner trackable
+      # has `_serialize_to_tensor`.
+      return trackable_has_serialize_to_tensor(obj._trackable)  # pylint: disable=protected-access
+    if t is trackable.Trackable:
+      # Base case. Return False since _serialize_to_tensors will raise a
+      # NotImplemented Error.
+      return False
+    elif "_serialize_to_tensors" in t.__dict__:
+      return True
+    elif "_gather_saveables_for_checkpoint" in t.__dict__:
+      return False
+  return False
+
+
+def _convert_to_string(x):
+  return compat.as_str(tensor_util.constant_value(x))
+
+
+class SaveableCompatibilityConverter(trackable.Trackable):
+  """Converts object's `SaveableObjects` to functions used in TF2 checkpointing.
+
+  A class that converts a Trackable object's `SaveableObjects` to save and
+  restore functions with the same signatures as
+  `Trackable._serialize_to_tensors` and `Trackable._restore_from_tensors`.
+  This class also produces a method for filling the object proto.
+  """
+
+  __slots__ = ("_obj", "_saveables")
+
+  def __init__(self, obj, saveables):
+    """Constructor.
+
+    Args:
+      obj: A Trackable object.
+      saveables: A list of saveables for `obj`.
+    """
+    self._obj = obj
+    self._saveables = saveables
+
+  @property
+  def obj(self):
+    return self._obj
+
+  @property
+  def saveables(self):
+    """Returns a list of SaveableObjects generated from the Trackable object."""
+    return self._saveables
+
+  def _serialize_to_tensors(self):
+    """Returns a dict of tensors to serialize."""
+    return saveable_object_to_tensor_dict(self.saveables)
+
+  def _restore_from_tensors(self, restored_tensors):
+    """Returns the restore ops defined in the Saveables."""
+    # Map restored tensors to the corresponding SaveableObjects, then call
+    # restore. There must be an exact match between restored tensors and the
+    # expected attributes.
+    expected_keys = []
+    for saveable in self.saveables:
+      expected_keys.extend(
+          trackable_utils.extract_local_name(_convert_to_string(spec.name))
+          for spec in saveable.specs)
+    if set(expected_keys) != restored_tensors.keys():
+      raise ValueError(f"Could not restore object {self._obj} because not all "
+                       "expected tensors were in the checkpoint."
+                       f"\n\tExpected: {expected_keys}"
+                       f"\n\tGot: {list(restored_tensors.keys())}")
+
+    return saveable_object_to_restore_fn(self.saveables)(restored_tensors)
+
+
+def saveable_object_to_tensor_dict(saveables):
+  """Converts a list of SaveableObjects to a tensor dictionary."""
+  tensor_dict = {}
+  for saveable in saveables:
+    for spec in saveable.specs:
+      name = _convert_to_string(spec.name)
+      slice_spec = _convert_to_string(spec.slice_spec)
+      # Currently, tensor dict cannot handle callable tensor values (which
+      # are needed for uninitialized variables), so keep using SaveSpec.
+      tensor = spec if callable(spec._tensor) else spec._tensor  # pylint: disable=protected-access
+      if slice_spec:
+        tensor_dict.setdefault(name, {})[slice_spec] = tensor
+      else:
+        tensor_dict[name] = tensor
+  return tensor_dict
+
+
+def saveable_object_to_restore_fn(saveables):
+  """Generates `Trackable._restore_from_tensors` from SaveableObjects."""
+
+  def _restore_from_tensors(restored_tensors):
+    restore_ops = {}
+
+    for saveable in saveables:
+      saveable_restored_tensors = []
+      for spec in saveable.specs:
+        name = trackable_utils.extract_local_name(_convert_to_string(spec.name))
+        slice_spec = _convert_to_string(spec.slice_spec)
+
+        maybe_tensor = restored_tensors[name]
+        if not isinstance(maybe_tensor, dict):
+          maybe_tensor = {"": maybe_tensor}
+
+        saveable_restored_tensors.append(maybe_tensor[slice_spec])
+      restore_ops[saveable.name] = saveable.restore(
+          saveable_restored_tensors, restored_shapes=None)
+    return restore_ops
+
+  return _restore_from_tensors
+
+
+def serialized_tensors_to_saveable_cache(serialized_tensors):
+  """Converts a tensor dict to a SaveableObject cache.
+
+  Args:
+    serialized_tensors: Map from Trackable to a tensor dict. The tensor dict
+      maps checkpoint key (-> slice_spec) -> Tensor
+
+  Returns:
+    A dict mapping Trackable objects to a map from local savable name to
+    SaveableObject.
+  """
+  saveables_cache = object_identity.ObjectIdentityWeakKeyDictionary()
+
+  for obj, tensor_dict in serialized_tensors.items():
+    if not tensor_dict: continue
+    if isinstance(obj, SaveableCompatibilityConverter):
+      trackable_obj = obj.obj
+      saveables_cache[trackable_obj] = {}
+      for saveable in obj.saveables:
+        local_name = trackable_utils.extract_local_name(saveable.name)
+        saveables_cache[trackable_obj][local_name] = [saveable]
+      continue
+
+    specs = []
+    # The local names and prefixes are computed to ensure that the generated
+    # SaveableObject can call `Trackable._restore_from_tensors()`
+    local_names = []
+    prefix = saveable_compat.get_saveable_name(obj) or ""
+    for checkpoint_key, maybe_tensor in tensor_dict.items():
+      # Make sure that `maybe_tensor` is a dict from `slice_spec` to `tensor`.
+      if not isinstance(maybe_tensor, dict):
+        maybe_tensor = {"": maybe_tensor}
+
+      for slice_spec, tensor in maybe_tensor.items():
+        if isinstance(tensor, saveable_object.SaveSpec):
+          specs.append(tensor)
+        else:
+          specs.append(saveable_object.SaveSpec(tensor,
+                                                slice_spec,
+                                                checkpoint_key))
+      local_names.append(trackable_utils.extract_local_name(checkpoint_key,
+                                                            prefix))
+
+    object_name = trackable_utils.extract_object_name(
+        next(iter(tensor_dict.keys())))
+    saveables_cache[obj] = {
+        trackable_utils.SERIALIZE_TO_TENSORS_NAME: [TrackableSaveable(
+            obj, specs, object_name, local_names=local_names, prefix=prefix)]}
+  return saveables_cache

tensorflow/python/training/saving/saveable_object_util_test.py

@@ -0,0 +1,251 @@
+# Copyright 2022 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.
+# ==============================================================================
+"""Tests for saveable_object_util."""
+
+import os
+
+from tensorflow.python.checkpoint import checkpoint
+from tensorflow.python.eager import context
+from tensorflow.python.eager import test
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import gen_resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.trackable import base
+from tensorflow.python.trackable import resource
+from tensorflow.python.training.saving import saveable_object
+from tensorflow.python.training.saving import saveable_object_util
+
+
+class _VarSaveable(saveable_object.SaveableObject):
+
+  def __init__(self, var, slice_spec, name):
+    specs = [saveable_object.SaveSpec(var.read_value(), slice_spec, name)]
+    super().__init__(var, specs, name)
+
+  def restore(self, restored_tensors, restored_shapes):
+    return self.op.assign(restored_tensors[0])
+
+
+def _create_converted_trackable(obj):
+  saveable_factories = saveable_object_util.saveable_objects_from_trackable(obj)
+  saveables = [factory(name) for name, factory in saveable_factories.items()]
+  return saveable_object_util.SaveableCompatibilityConverter(obj, saveables)
+
+
+class SaveableCompatibilityConverterTest(test.TestCase):
+
+  def test_convert_no_saveable(self):
+    t = base.Trackable()
+    converter = _create_converted_trackable(t)
+    self.assertEmpty(converter._serialize_to_tensors())
+    converter._restore_from_tensors({})
+
+    with self.assertRaisesRegex(ValueError, "Could not restore object"):
+      converter._restore_from_tensors({"": 0})
+
+  def test_convert_single_saveable(self):
+
+    class MyTrackable(base.Trackable):
+
+      def __init__(self):
+        self.a = variables.Variable(5.0)
+
+      def _gather_saveables_for_checkpoint(self):
+        return {"a": lambda name: _VarSaveable(self.a, "", name)}
+
+    t = MyTrackable()
+    converter = _create_converted_trackable(t)
+
+    serialized_tensors = converter._serialize_to_tensors()
+    self.assertLen(serialized_tensors, 1)
+    self.assertIn("a", serialized_tensors)
+    self.assertEqual(5, self.evaluate(serialized_tensors["a"]))
+
+    with self.assertRaisesRegex(ValueError, "Could not restore object"):
+      converter._restore_from_tensors({})
+    with self.assertRaisesRegex(ValueError, "Could not restore object"):
+      converter._restore_from_tensors({"not_a": 1.})
+
+    self.assertEqual(5, self.evaluate(t.a))
+    converter._restore_from_tensors({"a": 123.})
+    self.assertEqual(123, self.evaluate(t.a))
+
+  def test_convert_single_saveable_renamed(self):
+
+    class MyTrackable(base.Trackable):
+
+      def __init__(self):
+        self.a = variables.Variable(15.0)
+
+      def _gather_saveables_for_checkpoint(self):
+        return {"a": lambda name: _VarSaveable(self.a, "", name + "-value")}
+
+    t = MyTrackable()
+    converter = _create_converted_trackable(t)
+
+    serialized_tensors = converter._serialize_to_tensors()
+
+    self.assertLen(serialized_tensors, 1)
+    self.assertEqual(15, self.evaluate(serialized_tensors["a-value"]))
+
+    with self.assertRaisesRegex(ValueError, "Could not restore object"):
+      converter._restore_from_tensors({"a": 1.})
+
+    self.assertEqual(15, self.evaluate(t.a))
+    converter._restore_from_tensors({"a-value": 456.})
+    self.assertEqual(456, self.evaluate(t.a))
+
+
+class _MultiSpecSaveable(saveable_object.SaveableObject):
+
+  def __init__(self, obj, name):
+    self.obj = obj
+    specs = [
+        saveable_object.SaveSpec(obj.a, "", name + "-a"),
+        saveable_object.SaveSpec(obj.b, "", name + "-b")]
+    super(_MultiSpecSaveable, self).__init__(None, specs, name)
+
+  def restore(self, restored_tensors, restored_shapes):
+    del restored_shapes  # Unused.
+    self.obj.a.assign(restored_tensors[0])
+    self.obj.b.assign(restored_tensors[1])
+
+
+class State(resource.TrackableResource):
+
+  def __init__(self, initial_value):
+    super().__init__()
+    self._initial_value = initial_value
+    self._initialize()
+
+  def _create_resource(self):
+    return gen_resource_variable_ops.var_handle_op(
+        shape=[],
+        dtype=dtypes.float32,
+        shared_name=context.anonymous_name(),
+        name="StateVar",
+        container="")
+
+  def _initialize(self):
+    gen_resource_variable_ops.assign_variable_op(self.resource_handle,
+                                                 self._initial_value)
+
+  def _destroy_resource(self):
+    gen_resource_variable_ops.destroy_resource_op(self.resource_handle,
+                                                  ignore_lookup_error=True)
+
+  def read(self):
+    return gen_resource_variable_ops.read_variable_op(self.resource_handle,
+                                                      dtypes.float32)
+
+  def assign(self, value):
+    gen_resource_variable_ops.assign_variable_op(self.resource_handle, value)
+
+
+class _StateSaveable(saveable_object.SaveableObject):
+
+  def __init__(self, obj, name):
+    spec = saveable_object.SaveSpec(obj.read(), "", name)
+    self.obj = obj
+    super(_StateSaveable, self).__init__(obj, [spec], name)
+
+  def restore(self, restored_tensors, restored_shapes):
+    del restored_shapes  # Unused.
+    self.obj.assign(restored_tensors[0])
+
+
+class SaveableState(State):
+
+  def _gather_saveables_for_checkpoint(self):
+    return {
+        "value": lambda name: _StateSaveable(self, name)
+    }
+
+
+class TrackableState(State):
+
+  def _serialize_to_tensors(self):
+    return {
+        "value": self.read()
+    }
+
+  def _restore_from_tensors(self, restored_tensors):
+    return self.assign(restored_tensors["value"])
+
+
+class SaveableCompatibilityEndToEndTest(test.TestCase):
+
+  def test_checkpoint_comparison(self):
+    saveable_state = SaveableState(5.)
+    trackable_state = TrackableState(10.)
+
+    # First test that SaveableState and TrackableState are equivalent by
+    # saving a checkpoint with both objects and swapping values.
+
+    self.assertEqual(5, self.evaluate(saveable_state.read()))
+    self.assertEqual(10, self.evaluate(trackable_state.read()))
+
+    ckpt_path = os.path.join(self.get_temp_dir(), "ckpt")
+    checkpoint.Checkpoint(a=saveable_state, b=trackable_state).write(ckpt_path)
+
+    status = checkpoint.Checkpoint(b=saveable_state,
+                                   a=trackable_state).read(ckpt_path)
+    status.assert_consumed()
+
+    self.assertEqual(10, self.evaluate(saveable_state.read()))
+    self.assertEqual(5, self.evaluate(trackable_state.read()))
+
+    # Test that the converted SaveableState is compatible with the checkpoint
+    # saved above.
+    to_convert = SaveableState(0.0)
+
+    converted_saveable_state = _create_converted_trackable(to_convert)
+
+    checkpoint.Checkpoint(a=converted_saveable_state).read(
+        ckpt_path).assert_existing_objects_matched().expect_partial()
+    self.assertEqual(5, self.evaluate(to_convert.read()))
+
+    checkpoint.Checkpoint(b=converted_saveable_state).read(
+        ckpt_path).assert_existing_objects_matched().expect_partial()
+    self.assertEqual(10, self.evaluate(to_convert.read()))
+
+
+class HasSerializeToTensorTest(test.TestCase):
+  def test_has_serialize_to_tensor(self):
+    class ReturnsTrue(base.Trackable):
+      def _serialize_to_tensors(self):
+        return {}
+
+    class ReturnsFalse(base.Trackable):
+      pass
+
+    class SubclassReturnsFalse(ReturnsTrue):
+      def _gather_saveables_for_checkpoint(self):
+        return {}
+
+    self.assertTrue(saveable_object_util.trackable_has_serialize_to_tensor(
+        ReturnsTrue()))
+    self.assertFalse(saveable_object_util.trackable_has_serialize_to_tensor(
+        ReturnsFalse()))
+
+    # This should return False, because even though its parent class has
+    # `_serialize_to_tensors`, the class itself defines
+    # `_gather_saveables_for_checkpoint`.
+    self.assertFalse(saveable_object_util.trackable_has_serialize_to_tensor(
+        SubclassReturnsFalse()))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/saving/trace_saveable_util.py

@@ -0,0 +1,116 @@
+# 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.
+# ==============================================================================
+"""Utilities for tracing save and restore functions for SaveableObjects."""
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.framework import type_spec
+
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.training.saving import saveable_object
+from tensorflow.python.training.saving import saveable_object_util
+from tensorflow.python.util import nest
+
+
+def trace_save_restore_function_map(obj, factory_data_list):
+  """Traces all save and restore functions in the provided factory list.
+
+  Args:
+    obj: `Trackable` object.
+    factory_data_list: List of `_CheckpointFactoryData`.
+
+  Returns:
+    Dict mapping atttribute names to tuples of concrete save/restore functions.
+  """
+  saveable_fns = {}
+
+  for factory_data in factory_data_list:
+    saveable_factory = factory_data.factory
+    attribute_name = factory_data.name
+
+    # If object revives as a resource (or TPU/Mirrored) variable,
+    # there is no need to trace the save and restore functions.
+    if (resource_variable_ops.is_resource_variable(obj) or
+        resource_variable_ops.is_resource_variable(saveable_factory) or
+        not callable(saveable_factory)):
+      continue
+
+    concrete_save, concrete_restore = (
+        _trace_save_restore_functions(saveable_factory, obj))
+    if not concrete_save:
+      continue
+    saveable_fns[attribute_name] = (concrete_save, concrete_restore)
+  return saveable_fns
+
+
+def _trace_save_restore_functions(saveable_factory, obj):
+  """Traces save and restore functions."""
+  if saveable_object_util.is_factory_for_restored_saveable_object(
+      saveable_factory):
+    return (
+        saveable_factory.keywords["save_function"],
+        saveable_factory.keywords["restore_function"],
+    )
+
+  saveables = []  # Store the saveables in a data structure accessible to both
+  # the save and restore functions.
+
+  @def_function.function(
+      input_signature=[tensor_spec.TensorSpec([], dtypes.string)]
+  )
+  def save_fn(checkpoint_key):
+    maybe_saveable = saveable_factory(name=checkpoint_key)
+    if isinstance(maybe_saveable, saveable_object.SaveableObject):
+      maybe_saveable = [maybe_saveable]
+    saveables[:] = maybe_saveable
+
+    # Return list of all SaveSpecs created by the factory.
+    ret = []
+    for saveable in saveables:
+      for spec in saveable.specs:
+        ret.append({"name": spec.name, "tensor": spec.tensor,
+                    "slice_spec": spec.slice_spec})
+    return ret
+
+  concrete_save = save_fn.get_concrete_function()
+
+  # The SaveableObjects are produced when `save_fn` is traced.
+  saveables = saveable_object_util.validate_saveables_for_saved_model(
+      saveables, obj)
+  if not saveables:
+    return None, None
+
+  # Use the SaveSpecs to define the input signature of the restore function.
+  restored_type_specs = []
+  tensor_structure = []
+  for saveable in saveables:
+    saveable_tensor_structure = []
+    tensor_structure.append(saveable_tensor_structure)
+    for spec in saveable.specs:
+      restored_type_specs.append(type_spec.type_spec_from_value(spec.tensor))
+      saveable_tensor_structure.append(spec.name)
+
+  @def_function.function(input_signature=restored_type_specs)
+  def restore_fn(*restored_tensors):
+    structured_restored_tensors = nest.pack_sequence_as(
+        tensor_structure, restored_tensors)
+    for saveable, restored_tensors in zip(saveables,
+                                          structured_restored_tensors):
+      saveable.restore(restored_tensors, restored_shapes=None)
+    return 1  # Return dummy tensor
+
+  concrete_restore = restore_fn.get_concrete_function()
+  return concrete_save, concrete_restore

tensorflow/python/training/server_lib.py

@@ -0,0 +1,574 @@
+# 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.
+# ==============================================================================
+"""A Python interface for creating TensorFlow servers."""
+
+from tensorflow.core.protobuf import cluster_pb2
+from tensorflow.core.protobuf import device_filters_pb2
+from tensorflow.core.protobuf import tensorflow_server_pb2
+from tensorflow.python.client import pywrap_tf_session as c_api
+from tensorflow.python.framework import errors
+from tensorflow.python.util import compat
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+def _make_server_def(server_or_cluster_def, job_name, task_index, protocol,
+                     config):
+  """Creates a `tf.train.ServerDef` protocol buffer.
+
+  Args:
+    server_or_cluster_def: A `tf.train.ServerDef` or `tf.train.ClusterDef`
+      protocol buffer, or a `tf.train.ClusterSpec` object, describing the server
+      to be defined and/or the cluster of which it is a member.
+    job_name: (Optional.) Specifies the name of the job of which the server is a
+      member. Defaults to the value in `server_or_cluster_def`, if specified.
+    task_index: (Optional.) Specifies the task index of the server in its job.
+      Defaults to the value in `server_or_cluster_def`, if specified. Otherwise
+      defaults to 0 if the server's job has only one task.
+    protocol: (Optional.) Specifies the protocol to be used by the server.
+      Acceptable values include `"grpc", "grpc+verbs"`. Defaults to the value in
+      `server_or_cluster_def`, if specified. Otherwise defaults to `"grpc"`.
+    config: (Options.) A `tf.compat.v1.ConfigProto` that specifies default
+      configuration options for all sessions that run on this server.
+
+  Returns:
+    A `tf.train.ServerDef`.
+
+  Raises:
+    TypeError: If the arguments do not have the appropriate type.
+    ValueError: If an argument is not specified and cannot be inferred.
+  """
+  server_def = tensorflow_server_pb2.ServerDef()
+  if isinstance(server_or_cluster_def, tensorflow_server_pb2.ServerDef):
+    server_def.MergeFrom(server_or_cluster_def)
+    if job_name is not None:
+      server_def.job_name = job_name
+    if task_index is not None:
+      server_def.task_index = task_index
+    if protocol is not None:
+      server_def.protocol = protocol
+    if config is not None:
+      server_def.default_session_config.MergeFrom(config)
+  else:
+    try:
+      cluster_spec = ClusterSpec(server_or_cluster_def)
+    except TypeError:
+      raise TypeError("Could not convert `server_or_cluster_def` to a "
+                      "`tf.train.ServerDef` or `tf.train.ClusterSpec`.")
+    if job_name is None:
+      if len(cluster_spec.jobs) == 1:
+        job_name = cluster_spec.jobs[0]
+      else:
+        raise ValueError("Must specify an explicit `job_name`.")
+    if task_index is None:
+      task_indices = cluster_spec.task_indices(job_name)
+      if len(task_indices) == 1:
+        task_index = task_indices[0]
+      else:
+        raise ValueError("Must specify an explicit `task_index`.")
+    if protocol is None:
+      protocol = "grpc"
+
+    server_def = tensorflow_server_pb2.ServerDef(
+        cluster=cluster_spec.as_cluster_def(),
+        job_name=job_name,
+        task_index=task_index,
+        protocol=protocol)
+    if config is not None:
+      server_def.default_session_config.MergeFrom(config)
+  return server_def
+
+
+@tf_export("distribute.Server", v1=["distribute.Server", "train.Server"])
+@deprecation.deprecated_endpoints("train.Server")
+class Server:
+  """An in-process TensorFlow server, for use in distributed training.
+
+  A `tf.distribute.Server` instance encapsulates a set of devices and a
+  `tf.compat.v1.Session` target that
+  can participate in distributed training. A server belongs to a
+  cluster (specified by a `tf.train.ClusterSpec`), and
+  corresponds to a particular task in a named job. The server can
+  communicate with any other server in the same cluster.
+  """
+
+  def __init__(self,
+               server_or_cluster_def,
+               job_name=None,
+               task_index=None,
+               protocol=None,
+               config=None,
+               start=True):
+    """Creates a new server with the given definition.
+
+    The `job_name`, `task_index`, and `protocol` arguments are optional, and
+    override any information provided in `server_or_cluster_def`.
+
+    Args:
+      server_or_cluster_def: A `tf.train.ServerDef` or `tf.train.ClusterDef`
+        protocol buffer, or a `tf.train.ClusterSpec` object, describing the
+        server to be created and/or the cluster of which it is a member.
+      job_name: (Optional.) Specifies the name of the job of which the server is
+        a member. Defaults to the value in `server_or_cluster_def`, if
+        specified.
+      task_index: (Optional.) Specifies the task index of the server in its job.
+        Defaults to the value in `server_or_cluster_def`, if specified.
+        Otherwise defaults to 0 if the server's job has only one task.
+      protocol: (Optional.) Specifies the protocol to be used by the server.
+        Acceptable values include `"grpc", "grpc+verbs"`. Defaults to the value
+        in `server_or_cluster_def`, if specified. Otherwise defaults to
+        `"grpc"`.
+      config: (Options.) A `tf.compat.v1.ConfigProto` that specifies default
+        configuration options for all sessions that run on this server.
+      start: (Optional.) Boolean, indicating whether to start the server after
+        creating it. Defaults to `True`.
+
+    Raises:
+      tf.errors.OpError: Or one of its subclasses if an error occurs while
+        creating the TensorFlow server.
+    """
+    self._server_def = _make_server_def(server_or_cluster_def, job_name,
+                                        task_index, protocol, config)
+    self._server = c_api.TF_NewServer(self._server_def.SerializeToString())
+    if start:
+      self.start()
+
+  def __del__(self):
+    # At shutdown, `errors` may have been garbage collected.
+    if errors is not None:
+      exception = errors.UnimplementedError
+    else:
+      exception = Exception
+    try:
+      c_api.TF_ServerStop(self._server)
+      # Clean shutdown of servers is not yet implemented, so
+      # we leak instead of calling c_api.TF_DeleteServer here.
+      # See:
+      # https://github.com/tensorflow/tensorflow/blob/0495317a6e9dd4cac577b9d5cf9525e62b571018/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.h#L73
+    except AttributeError:
+      # At shutdown, `c_api` may have been garbage collected.
+      pass
+    except exception:
+      pass
+    self._server = None
+
+  def start(self):
+    """Starts this server.
+
+    Raises:
+      tf.errors.OpError: Or one of its subclasses if an error occurs while
+        starting the TensorFlow server.
+    """
+    c_api.TF_ServerStart(self._server)
+
+  def join(self):
+    """Blocks until the server has shut down.
+
+    This method currently blocks forever.
+
+    Raises:
+      tf.errors.OpError: Or one of its subclasses if an error occurs while
+        joining the TensorFlow server.
+    """
+    c_api.TF_ServerJoin(self._server)
+
+  @property
+  def server_def(self):
+    """Returns the `tf.train.ServerDef` for this server.
+
+    Returns:
+      A `tf.train.ServerDef` protocol buffer that describes the configuration
+      of this server.
+    """
+    return self._server_def
+
+  @property
+  def target(self):
+    """Returns the target for a `tf.compat.v1.Session` to connect to this server.
+
+    To create a
+    `tf.compat.v1.Session` that
+    connects to this server, use the following snippet:
+
+    ```python
+    server = tf.distribute.Server(...)
+    with tf.compat.v1.Session(server.target):
+      # ...
+    ```
+
+    Returns:
+      A string containing a session target for this server.
+    """
+    return c_api.TF_ServerTarget(self._server)
+
+  @staticmethod
+  def create_local_server(config=None, start=True):
+    """Creates a new single-process cluster running on the local host.
+
+    This method is a convenience wrapper for creating a
+    `tf.distribute.Server` with a `tf.train.ServerDef` that specifies a
+    single-process cluster containing a single task in a job called
+    `"local"`.
+
+    Args:
+      config: (Options.) A `tf.compat.v1.ConfigProto` that specifies default
+        configuration options for all sessions that run on this server.
+      start: (Optional.) Boolean, indicating whether to start the server after
+        creating it. Defaults to `True`.
+
+    Returns:
+      A local `tf.distribute.Server`.
+    """
+    # Specifying port 0 means that the OS will choose a free port for the
+    # server.
+    return Server({"localhost": ["localhost:0"]},
+                  protocol="grpc",
+                  config=config,
+                  start=start)
+
+
+@tf_export("train.ClusterSpec")
+class ClusterSpec:
+  """Represents a cluster as a set of "tasks", organized into "jobs".
+
+  A `tf.train.ClusterSpec` represents the set of processes that
+  participate in a distributed TensorFlow computation. Every
+  `tf.distribute.Server` is constructed in a particular cluster.
+
+  To create a cluster with two jobs and five tasks, you specify the
+  mapping from job names to lists of network addresses (typically
+  hostname-port pairs).
+
+  ```python
+  cluster = tf.train.ClusterSpec({"worker": ["worker0.example.com:2222",
+                                             "worker1.example.com:2222",
+                                             "worker2.example.com:2222"],
+                                  "ps": ["ps0.example.com:2222",
+                                         "ps1.example.com:2222"]})
+  ```
+
+  Each job may also be specified as a sparse mapping from task indices
+  to network addresses. This enables a server to be configured without
+  needing to know the identity of (for example) all other worker
+  tasks:
+
+  ```python
+  cluster = tf.train.ClusterSpec({"worker": {1: "worker1.example.com:2222"},
+                                  "ps": ["ps0.example.com:2222",
+                                         "ps1.example.com:2222"]})
+  ```
+  """
+
+  def __init__(self, cluster):
+    """Creates a `ClusterSpec`.
+
+    Args:
+      cluster: A dictionary mapping one or more job names to (i) a list of
+        network addresses, or (ii) a dictionary mapping integer task indices to
+        network addresses; or a `tf.train.ClusterDef` protocol buffer.
+
+    Raises:
+      TypeError: If `cluster` is not a dictionary mapping strings to lists
+        of strings, and not a `tf.train.ClusterDef` protobuf.
+    """
+    if isinstance(cluster, dict):
+      self._cluster_spec = {}
+      for job_name, tasks in cluster.items():
+        if isinstance(tasks, (list, tuple)):
+          job_tasks = {i: task for i, task in enumerate(tasks)}
+        elif isinstance(tasks, dict):
+          job_tasks = {int(i): task for i, task in tasks.items()}
+        else:
+          raise TypeError("The tasks for job %r must be a list or a dictionary "
+                          "from integers to strings." % job_name)
+        self._cluster_spec[job_name] = job_tasks
+      self._make_cluster_def()
+    elif isinstance(cluster, cluster_pb2.ClusterDef):
+      self._cluster_def = cluster
+      self._cluster_spec = {}
+      for job_def in self._cluster_def.job:
+        self._cluster_spec[job_def.name] = {
+            i: t for i, t in job_def.tasks.items()
+        }
+    elif isinstance(cluster, ClusterSpec):
+      self._cluster_def = cluster_pb2.ClusterDef()
+      self._cluster_def.MergeFrom(cluster.as_cluster_def())
+      self._cluster_spec = {}
+      for job_def in self._cluster_def.job:
+        self._cluster_spec[job_def.name] = {
+            i: t for i, t in job_def.tasks.items()
+        }
+    else:
+      raise TypeError("`cluster` must be a dictionary mapping one or more "
+                      "job names to lists of network addresses, or a "
+                      "`ClusterDef` protocol buffer")
+
+  def __bool__(self):
+    return bool(self._cluster_spec)
+
+  # Python 2.x
+  __nonzero__ = __bool__
+
+  def __eq__(self, other):
+    return self._cluster_spec == other
+
+  def __ne__(self, other):
+    return self._cluster_spec != other
+
+  def __repr__(self):
+    key_values = self.as_dict()
+    string_items = [
+        repr(k) + ": " + repr(key_values[k]) for k in sorted(key_values)
+    ]
+    return "ClusterSpec({" + ", ".join(string_items) + "})"
+
+  def as_dict(self):
+    """Returns a dictionary from job names to their tasks.
+
+    For each job, if the task index space is dense, the corresponding
+    value will be a list of network addresses; otherwise it will be a
+    dictionary mapping (sparse) task indices to the corresponding
+    addresses.
+
+    Returns:
+      A dictionary mapping job names to lists or dictionaries
+      describing the tasks in those jobs.
+    """
+    ret = {}
+    for job in self.jobs:
+      task_indices = self.task_indices(job)
+      if len(task_indices) == 0:
+        ret[job] = {}
+        continue
+      if max(task_indices) + 1 == len(task_indices):
+        # Return a list because the task indices are dense. This
+        # matches the behavior of `as_dict()` before support for
+        # sparse jobs was added.
+        ret[job] = self.job_tasks(job)
+      else:
+        ret[job] = {i: self.task_address(job, i) for i in task_indices}
+    return ret
+
+  def as_cluster_def(self):
+    """Returns a `tf.train.ClusterDef` protocol buffer based on this cluster."""
+    return self._cluster_def
+
+  @property
+  def jobs(self):
+    """Returns a list of job names in this cluster.
+
+    Returns:
+      A list of strings, corresponding to the names of jobs in this cluster.
+    """
+    return list(self._cluster_spec.keys())
+
+  def num_tasks(self, job_name):
+    """Returns the number of tasks defined in the given job.
+
+    Args:
+      job_name: The string name of a job in this cluster.
+
+    Returns:
+      The number of tasks defined in the given job.
+
+    Raises:
+      ValueError: If `job_name` does not name a job in this cluster.
+    """
+    try:
+      job = self._cluster_spec[job_name]
+    except KeyError:
+      raise ValueError("No such job in cluster: %r" % job_name)
+    return len(job)
+
+  def task_indices(self, job_name):
+    """Returns a list of valid task indices in the given job.
+
+    Args:
+      job_name: The string name of a job in this cluster.
+
+    Returns:
+      A list of valid task indices in the given job.
+
+    Raises:
+      ValueError: If `job_name` does not name a job in this cluster,
+      or no task with index `task_index` is defined in that job.
+    """
+    try:
+      job = self._cluster_spec[job_name]
+    except KeyError:
+      raise ValueError("No such job in cluster: %r" % job_name)
+    return list(sorted(job.keys()))
+
+  def task_address(self, job_name, task_index):
+    """Returns the address of the given task in the given job.
+
+    Args:
+      job_name: The string name of a job in this cluster.
+      task_index: A non-negative integer.
+
+    Returns:
+      The address of the given task in the given job.
+
+    Raises:
+      ValueError: If `job_name` does not name a job in this cluster,
+      or no task with index `task_index` is defined in that job.
+    """
+    try:
+      job = self._cluster_spec[job_name]
+    except KeyError:
+      raise ValueError("No such job in cluster: %r" % job_name)
+    try:
+      return job[task_index]
+    except KeyError:
+      raise ValueError("No task with index %r in job %r" %
+                       (task_index, job_name))
+
+  def job_tasks(self, job_name):
+    """Returns a mapping from task ID to address in the given job.
+
+    NOTE: For backwards compatibility, this method returns a list. If
+    the given job was defined with a sparse set of task indices, the
+    length of this list may not reflect the number of tasks defined in
+    this job. Use the `tf.train.ClusterSpec.num_tasks` method
+    to find the number of tasks defined in a particular job.
+
+    Args:
+      job_name: The string name of a job in this cluster.
+
+    Returns:
+      A list of task addresses, where the index in the list
+      corresponds to the task index of each task. The list may contain
+      `None` if the job was defined with a sparse set of task indices.
+
+    Raises:
+      ValueError: If `job_name` does not name a job in this cluster.
+    """
+    try:
+      job = self._cluster_spec[job_name]
+    except KeyError:
+      raise ValueError("No such job in cluster: %r" % job_name)
+    ret = [None for _ in range(max(job.keys()) + 1)]
+    for i, task in job.items():
+      ret[i] = task
+    return ret
+
+  def _make_cluster_def(self):
+    """Creates a `tf.train.ClusterDef` based on the given `cluster_spec`.
+
+    Raises:
+      TypeError: If `cluster_spec` is not a dictionary mapping strings to lists
+        of strings.
+    """
+    self._cluster_def = cluster_pb2.ClusterDef()
+
+    # NOTE(mrry): Sort by job_name to produce deterministic protobufs.
+    for job_name, tasks in sorted(self._cluster_spec.items()):
+      try:
+        job_name = compat.as_bytes(job_name)
+      except TypeError:
+        raise TypeError("Job name %r must be bytes or unicode" % job_name)
+
+      job_def = self._cluster_def.job.add()
+      job_def.name = job_name
+
+      for i, task_address in sorted(tasks.items()):
+        try:
+          task_address = compat.as_bytes(task_address)
+        except TypeError:
+          raise TypeError("Task address %r must be bytes or unicode" %
+                          task_address)
+        job_def.tasks[i] = task_address
+
+
+@tf_export("config.experimental.ClusterDeviceFilters")
+class ClusterDeviceFilters:
+  """Represent a collection of device filters for the remote workers in cluster.
+
+  NOTE: this is an experimental API and subject to changes.
+
+  Set device filters for selective jobs and tasks. For each remote worker, the
+  device filters are a list of strings. When any filters are present, the remote
+  worker will ignore all devices which do not match any of its filters. Each
+  filter can be partially specified, e.g. "/job:ps", "/job:worker/replica:3",
+  etc. Note that a device is always visible to the worker it is located on.
+
+  For example, to set the device filters for a parameter server cluster:
+
+  ```python
+  cdf = tf.config.experimental.ClusterDeviceFilters()
+  for i in range(num_workers):
+    cdf.set_device_filters('worker', i, ['/job:ps'])
+  for i in range(num_ps):
+    cdf.set_device_filters('ps', i, ['/job:worker'])
+
+  tf.config.experimental_connect_to_cluster(cluster_def,
+                                            cluster_device_filters=cdf)
+  ```
+
+  The device filters can be partically specified. For remote tasks that do not
+  have device filters specified, all devices will be visible to them.
+  """
+
+  def __init__(self):
+    # `_device_filters` is a dict mapping job names to job device filters.
+    # Job device filters further maps task IDs to task device filters.
+    # Task device filters are a list of strings, each one is a device filter.
+    self._device_filters = {}
+
+    # Serialized protobuf for cluster device filters.
+    self._cluster_device_filters = None
+
+  def set_device_filters(self, job_name, task_index, device_filters):
+    """Set the device filters for given job name and task id."""
+    assert all(isinstance(df, str) for df in device_filters)
+    self._device_filters.setdefault(job_name, {})
+    self._device_filters[job_name][task_index] = [df for df in device_filters]
+    # Due to updates in data, invalidate the serialized proto cache.
+    self._cluster_device_filters = None
+
+  def _as_cluster_device_filters(self):
+    """Returns a serialized protobuf of cluster device filters."""
+    if self._cluster_device_filters:
+      return self._cluster_device_filters
+
+    self._make_cluster_device_filters()
+    return self._cluster_device_filters
+
+  def _make_cluster_device_filters(self):
+    """Creates `ClusterDeviceFilters` proto based on the `_device_filters`.
+
+    Raises:
+      TypeError: If `_device_filters` is not a dictionary mapping strings to
+      a map of task indices and device filters.
+    """
+    self._cluster_device_filters = device_filters_pb2.ClusterDeviceFilters()
+
+    # Sort by job_name to produce deterministic protobufs.
+    for job_name, tasks in sorted(self._device_filters.items()):
+      try:
+        job_name = compat.as_bytes(job_name)
+      except TypeError:
+        raise TypeError("Job name %r must be bytes or unicode" % job_name)
+
+      jdf = self._cluster_device_filters.jobs.add()
+      jdf.name = job_name
+
+      for i, task_device_filters in sorted(tasks.items()):
+        for tdf in task_device_filters:
+          try:
+            tdf = compat.as_bytes(tdf)
+          except TypeError:
+            raise TypeError("Device filter %r must be bytes or unicode" % tdf)
+          jdf.tasks[i].device_filters.append(tdf)

tensorflow/python/training/server_lib_multiple_containers_test.py

@@ -0,0 +1,60 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Tests for tf.GrpcServer."""
+
+from tensorflow.python.client import session
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import server_lib
+
+
+class MultipleContainersTest(test.TestCase):
+
+  # Verifies behavior of tf.Session.reset() with multiple containers using
+  # tf.container.
+  # TODO(b/34465411): Starting multiple servers with different configurations
+  # in the same test is flaky. Move this test case back into
+  # "server_lib_test.py" when this is no longer the case.
+  @test_util.run_deprecated_v1
+  def testMultipleContainers(self):
+    with ops.container("test0"):
+      v0 = variables.Variable(1.0, name="v0")
+    with ops.container("test1"):
+      v1 = variables.Variable(2.0, name="v0")
+    server = server_lib.Server.create_local_server()
+    sess = session.Session(server.target)
+    sess.run(variables.global_variables_initializer())
+    self.assertAllEqual(1.0, sess.run(v0))
+    self.assertAllEqual(2.0, sess.run(v1))
+
+    # Resets container. Session aborts.
+    session.Session.reset(server.target, ["test0"])
+    with self.assertRaises(errors_impl.AbortedError):
+      sess.run(v1)
+
+    # Connects to the same target. Device memory for the v0 would have
+    # been released, so it will be uninitialized. But v1 should still
+    # be valid.
+    sess = session.Session(server.target)
+    with self.assertRaises(errors_impl.FailedPreconditionError):
+      sess.run(v0)
+    self.assertAllEqual(2.0, sess.run(v1))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/server_lib_same_variables_clear_container_test.py

@@ -0,0 +1,79 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Tests for tf.GrpcServer."""
+
+from tensorflow.python.client import session
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import server_lib
+
+
+class SameVariablesClearContainerTest(test.TestCase):
+
+  # Verifies behavior of tf.Session.reset() with multiple containers using
+  # default container names as defined by the target name.
+  # TODO(b/34465411): Starting multiple servers with different configurations
+  # in the same test is flaky. Move this test case back into
+  # "server_lib_test.py" when this is no longer the case.
+  def testSameVariablesClearContainer(self):
+    # Starts two servers with different names so they map to different
+    # resource "containers".
+    server0 = server_lib.Server(
+        {
+            "local0": ["localhost:0"]
+        }, protocol="grpc", start=True)
+    server1 = server_lib.Server(
+        {
+            "local1": ["localhost:0"]
+        }, protocol="grpc", start=True)
+
+    # Creates a graph with 2 variables.
+    with ops.Graph().as_default():
+      v0 = variables.Variable(1.0, name="v0")
+      v1 = variables.Variable(2.0, name="v0")
+
+      # Initializes the variables. Verifies that the values are correct.
+      sess_0 = session.Session(server0.target)
+      sess_1 = session.Session(server1.target)
+      sess_0.run(v0.initializer)
+      sess_1.run(v1.initializer)
+      self.assertAllEqual(1.0, sess_0.run(v0))
+      self.assertAllEqual(2.0, sess_1.run(v1))
+
+      # Resets container "local0". Verifies that v0 is no longer initialized.
+      session.Session.reset(server0.target, ["local0"])
+      _ = session.Session(server0.target)
+      with self.assertRaises(errors_impl.FailedPreconditionError):
+        self.evaluate(v0)
+      # Reinitializes v0 for the following test.
+      self.evaluate(v0.initializer)
+
+      # Verifies that v1 is still valid.
+      self.assertAllEqual(2.0, sess_1.run(v1))
+
+      # Resets container "local1". Verifies that v1 is no longer initialized.
+      session.Session.reset(server1.target, ["local1"])
+      _ = session.Session(server1.target)
+      with self.assertRaises(errors_impl.FailedPreconditionError):
+        self.evaluate(v1)
+      # Verifies that v0 is still valid.
+      _ = session.Session(server0.target)
+      self.assertAllEqual(1.0, self.evaluate(v0))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/server_lib_same_variables_clear_test.py

@@ -0,0 +1,66 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Tests for tf.GrpcServer."""
+
+from tensorflow.python.client import session
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import server_lib
+
+
+class SameVariablesClearTest(test.TestCase):
+
+  # Verifies behavior of tf.Session.reset().
+  # TODO(b/34465411): Starting multiple servers with different configurations
+  # in the same test is flaky. Move this test case back into
+  # "server_lib_test.py" when this is no longer the case.
+  @test_util.run_deprecated_v1
+  def testSameVariablesClear(self):
+    server = server_lib.Server.create_local_server()
+
+    # Creates a graph with 2 variables.
+    v0 = variables.Variable([[2, 1]], name="v0")
+    v1 = variables.Variable([[1], [2]], name="v1")
+    v2 = math_ops.matmul(v0, v1)
+
+    # Verifies that both sessions connecting to the same target return
+    # the same results.
+    sess_1 = session.Session(server.target)
+    sess_2 = session.Session(server.target)
+    sess_1.run(variables.global_variables_initializer())
+    self.assertAllEqual([[4]], sess_1.run(v2))
+    self.assertAllEqual([[4]], sess_2.run(v2))
+
+    # Resets target. sessions abort. Use sess_2 to verify.
+    session.Session.reset(server.target)
+    with self.assertRaises(errors_impl.AbortedError):
+      self.assertAllEqual([[4]], sess_2.run(v2))
+
+    # Connects to the same target. Device memory for the variables would have
+    # been released, so they will be uninitialized.
+    sess_2 = session.Session(server.target)
+    with self.assertRaises(errors_impl.FailedPreconditionError):
+      sess_2.run(v2)
+    # Reinitializes the variables.
+    sess_2.run(variables.global_variables_initializer())
+    self.assertAllEqual([[4]], sess_2.run(v2))
+    sess_2.close()
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/server_lib_same_variables_no_clear_test.py

@@ -0,0 +1,53 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Tests for tf.GrpcServer."""
+
+from tensorflow.python.client import session
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.platform import test
+from tensorflow.python.training import server_lib
+
+
+class SameVariablesNoClearTest(test.TestCase):
+
+  # Verifies behavior of multiple variables with multiple sessions connecting to
+  # the same server.
+  # TODO(b/34465411): Starting multiple servers with different configurations
+  # in the same test is flaky. Move this test case back into
+  # "server_lib_test.py" when this is no longer the case.
+  @test_util.run_v1_only(
+      "This exercises tensor lookup via names which is not supported in V2.")
+  def testSameVariablesNoClear(self):
+    server = server_lib.Server.create_local_server()
+
+    with session.Session(server.target) as sess_1:
+      v0 = variable_v1.VariableV1([[2, 1]], name="v0")
+      v1 = variable_v1.VariableV1([[1], [2]], name="v1")
+      v2 = math_ops.matmul(v0, v1)
+      sess_1.run([v0.initializer, v1.initializer])
+      self.assertAllEqual([[4]], sess_1.run(v2))
+
+    with session.Session(server.target) as sess_2:
+      new_v0 = ops.get_default_graph().get_tensor_by_name("v0:0")
+      new_v1 = ops.get_default_graph().get_tensor_by_name("v1:0")
+      new_v2 = math_ops.matmul(new_v0, new_v1)
+      self.assertAllEqual([[4]], sess_2.run(new_v2))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/server_lib_sparse_job_test.py

@@ -0,0 +1,41 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Tests for tf.GrpcServer."""
+
+from tensorflow.python.client import session
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import test
+from tensorflow.python.training import server_lib
+
+
+class SparseJobTest(test.TestCase):
+
+  # TODO(b/34465411): Starting multiple servers with different configurations
+  # in the same test is flaky. Move this test case back into
+  # "server_lib_test.py" when this is no longer the case.
+  @test_util.run_deprecated_v1
+  def testSparseJob(self):
+    server = server_lib.Server({"local": {37: "localhost:0"}})
+    with ops.device("/job:local/task:37"):
+      a = constant_op.constant(1.0)
+
+    with session.Session(server.target) as sess:
+      self.assertEqual(1.0, self.evaluate(a))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/server_lib_test.py

@@ -0,0 +1,677 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Tests for tf.GrpcServer."""
+
+import time
+
+import numpy as np
+
+from xla.tsl.protobuf import rpc_options_pb2
+from tensorflow.core.protobuf import cluster_pb2
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.core.protobuf import tensorflow_server_pb2
+from tensorflow.python.client import session
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import data_flow_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import input as input_ops
+from tensorflow.python.training import queue_runner_impl
+from tensorflow.python.training import server_lib
+
+
+class GrpcServerTest(test.TestCase):
+
+  def __init__(self, methodName="runTest"):  # pylint: disable=invalid-name
+    super(GrpcServerTest, self).__init__(methodName)
+    self._cached_server = server_lib.Server.create_local_server()
+
+  def testRunStep(self):
+    server = self._cached_server
+    with ops.Graph().as_default():
+      with session.Session(server.target) as sess:
+        c = constant_op.constant([[2, 1]])
+        d = constant_op.constant([[1], [2]])
+        e = math_ops.matmul(c, d)
+        self.assertAllEqual([[4]], sess.run(e))
+      # TODO(mrry): Add `server.stop()` and `server.join()` when these work.
+
+  def testMultipleSessions(self):
+    server = self._cached_server
+    with ops.Graph().as_default():
+      c = constant_op.constant([[2, 1]])
+      d = constant_op.constant([[1], [2]])
+      e = math_ops.matmul(c, d)
+
+      sess_1 = session.Session(server.target)
+      sess_2 = session.Session(server.target)
+
+      self.assertAllEqual([[4]], sess_1.run(e))
+      self.assertAllEqual([[4]], sess_2.run(e))
+
+      sess_1.close()
+      sess_2.close()
+      # TODO(mrry): Add `server.stop()` and `server.join()` when these work.
+
+  # Verifies various reset failures.
+  def testResetFails(self):
+    with ops.Graph().as_default():
+      # Creates variable with container name.
+      with ops.container("test0"):
+        v0 = variable_v1.VariableV1(1.0, name="v0")
+      # Creates variable with default container.
+      v1 = variable_v1.VariableV1(2.0, name="v1")
+      # Verifies resetting the non-existent target returns error.
+      with self.assertRaises(errors_impl.NotFoundError):
+        session.Session.reset("nonexistent", ["test0"])
+
+      # Verifies resetting with config.
+      # Verifies that resetting target with no server times out.
+      with self.assertRaises(errors_impl.DeadlineExceededError):
+        session.Session.reset(
+            "grpc://localhost:0",
+            ["test0"],
+            config=config_pb2.ConfigProto(operation_timeout_in_ms=5),
+        )
+
+      # Verifies no containers are reset with non-existent container.
+      server = self._cached_server
+      sess = session.Session(server.target)
+      sess.run(variables.global_variables_initializer())
+      self.assertAllEqual(1.0, sess.run(v0))
+      self.assertAllEqual(2.0, sess.run(v1))
+      # No container is reset, but the server is reset.
+      session.Session.reset(server.target, ["test1"])
+      # Verifies that both variables are still valid.
+      sess = session.Session(server.target)
+      self.assertAllEqual(1.0, sess.run(v0))
+      self.assertAllEqual(2.0, sess.run(v1))
+
+  def _useRPCConfig(self):
+    """Return a `tf.compat.v1.ConfigProto` that ensures we use the RPC stack for tests.
+
+    This configuration ensures that we continue to exercise the gRPC
+    stack when testing, rather than using the in-process optimization,
+    which avoids using gRPC as the transport between a client and
+    master in the same process.
+
+    Returns:
+      A `tf.compat.v1.ConfigProto`.
+    """
+    return config_pb2.ConfigProto(
+        rpc_options=rpc_options_pb2.RPCOptions(
+            use_rpc_for_inprocess_master=True
+        )
+    )
+
+  def testLargeConstant(self):
+    server = self._cached_server
+    with session.Session(server.target, config=self._useRPCConfig()) as sess:
+      const_val = np.empty([10000, 3000], dtype=np.float32)
+      const_val.fill(0.5)
+      c = constant_op.constant(const_val)
+      shape_t = array_ops.shape(c)
+      self.assertAllEqual([10000, 3000], sess.run(shape_t))
+
+  def testLargeFetch(self):
+    server = self._cached_server
+    with session.Session(server.target, config=self._useRPCConfig()) as sess:
+      c = array_ops.fill([10000, 3000], 0.5)
+      expected_val = np.empty([10000, 3000], dtype=np.float32)
+      expected_val.fill(0.5)
+      self.assertAllEqual(expected_val, sess.run(c))
+
+  def testLargeFeed(self):
+    server = self._cached_server
+    with session.Session(server.target, config=self._useRPCConfig()) as sess:
+      feed_val = np.empty([10000, 3000], dtype=np.float32)
+      feed_val.fill(0.5)
+      p = array_ops.placeholder(dtypes.float32, shape=[10000, 3000])
+      min_t = math_ops.reduce_min(p)
+      max_t = math_ops.reduce_max(p)
+      min_val, max_val = sess.run([min_t, max_t], feed_dict={p: feed_val})
+      self.assertEqual(0.5, min_val)
+      self.assertEqual(0.5, max_val)
+
+  def testCloseCancelsBlockingOperation(self):
+    server = self._cached_server
+    with ops.Graph().as_default():
+      sess = session.Session(server.target, config=self._useRPCConfig())
+
+      q = data_flow_ops.FIFOQueue(10, [dtypes.float32])
+      enqueue_op = q.enqueue(37.0)
+      dequeue_t = q.dequeue()
+
+      sess.run(enqueue_op)
+      sess.run(dequeue_t)
+
+      def blocking_dequeue():
+        with self.assertRaisesRegex(
+            errors_impl.CancelledError, "Session::Close"
+        ):
+          sess.run(dequeue_t)
+
+      blocking_thread = self.checkedThread(blocking_dequeue)
+      blocking_thread.start()
+      time.sleep(0.5)
+      sess.close()
+      blocking_thread.join()
+
+  def testInteractiveSession(self):
+    server = self._cached_server
+    # Session creation will warn (in C++) that the place_pruned_graph option
+    # is not supported, but it should successfully ignore it.
+    sess = session.InteractiveSession(server.target)
+    c = constant_op.constant(42.0)
+    self.assertEqual(42.0, self.evaluate(c))
+    sess.close()
+
+  def testSetConfiguration(self):
+    config = config_pb2.ConfigProto(
+        gpu_options=config_pb2.GPUOptions(per_process_gpu_memory_fraction=0.1)
+    )
+
+    # Configure a server using the default local server options.
+    server = server_lib.Server.create_local_server(config=config, start=False)
+    self.assertEqual(
+        0.1,
+        server.server_def.default_session_config.gpu_options.per_process_gpu_memory_fraction,
+    )
+
+    # Configure a server using an explicit ServerDefd with an
+    # overridden config.
+    cluster_def = server_lib.ClusterSpec(
+        {"localhost": ["localhost:0"]}
+    ).as_cluster_def()
+    server_def = tensorflow_server_pb2.ServerDef(
+        cluster=cluster_def, job_name="localhost", task_index=0, protocol="grpc"
+    )
+    server = server_lib.Server(server_def, config=config, start=False)
+    self.assertEqual(
+        0.1,
+        server.server_def.default_session_config.gpu_options.per_process_gpu_memory_fraction,
+    )
+
+  def testRestartedMaster(self):
+    master_old = server_lib.Server.create_local_server()
+    master_new = server_lib.Server.create_local_server()
+    worker = self._cached_server
+
+    def get_cluster_def(master, worker):
+      cluster_def = cluster_pb2.ClusterDef()
+      job = cluster_def.job.add()
+      job.name = "master"
+      job.tasks[0] = master.target[len("grpc://") :]
+      job = cluster_def.job.add()
+      job.name = "worker"
+      job.tasks[0] = worker.target[len("grpc://") :]
+      return cluster_def
+
+    def check_session_devices(sess):
+      # Make sure we have the correct set of cluster devices
+      devices = sess.list_devices()
+      device_names = set(d.name for d in devices)
+      self.assertIn("/job:master/replica:0/task:0/device:CPU:0", device_names)
+      self.assertIn("/job:worker/replica:0/task:0/device:CPU:0", device_names)
+
+    with ops.Graph().as_default():
+      # Construct a simple graph that runs ops on remote worker
+      with ops.device("/job:worker/replica:0/task:0/device:CPU:0"):
+        a = constant_op.constant([1.0])
+        b = a + a
+
+      config = config_pb2.ConfigProto(
+          cluster_def=get_cluster_def(master_old, worker)
+      )
+      sess_old = session.Session(master_old.target, config=config)
+      check_session_devices(sess_old)
+
+      # Create a session with the new master and the worker.
+      # The new master has the same task name ('/job:master/replica:0/task:0')
+      # as the old master, but is initiated from a different server thus has a
+      # different incarnation. This triggers the WorkerSession on worker with
+      # the old master incarnation to be garbage collected.
+
+      config = config_pb2.ConfigProto(
+          cluster_def=get_cluster_def(master_new, worker)
+      )
+      sess_new = session.Session(master_new.target, config=config)
+      check_session_devices(sess_new)
+
+      # Running on worker with the new session should work as expected
+      v = sess_new.run(b)
+      self.assertAllEqual(v, [2.0])
+
+      # Running on worker with the old session should raise an exception since
+      # the WorkerSession of the old session has been garbage collected
+      with self.assertRaisesRegex(
+          errors_impl.AbortedError, "Session handle is not found"
+      ):
+        sess_old.run(b)
+
+    sess_old.close()
+    sess_new.close()
+
+  def testInvalidHostname(self):
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError, "port"):
+      _ = server_lib.Server(
+          {"local": ["localhost"]}, job_name="local", task_index=0
+      )
+
+  def testTimeoutRaisesException(self):
+    server = self._cached_server
+    with ops.Graph().as_default():
+      q = data_flow_ops.FIFOQueue(1, [dtypes.float32])
+      blocking_t = q.dequeue()
+
+      with session.Session(server.target) as sess:
+        with self.assertRaises(errors_impl.DeadlineExceededError):
+          sess.run(
+              blocking_t, options=config_pb2.RunOptions(timeout_in_ms=1000)
+          )
+
+      with session.Session(server.target, config=self._useRPCConfig()) as sess:
+        with self.assertRaises(errors_impl.DeadlineExceededError):
+          sess.run(
+              blocking_t, options=config_pb2.RunOptions(timeout_in_ms=1000)
+          )
+
+  def testTwoServersSamePort(self):
+    # Starting a server with the same target as the cached server should fail.
+    server = self._cached_server
+    with self.assertRaises(errors_impl.UnknownError):
+      _ = server_lib.Server({"local_2": [server.target[len("grpc://") :]]})
+
+  def testExtendAfterQueueRunners(self):
+    server = self._cached_server
+    with session.Session(server.target) as sess:
+      input_queue = input_ops.input_producer(
+          constant_op.constant([0.0], dtype=dtypes.float32)
+      )
+      self.assertIsNotNone(input_queue)
+
+      var = variable_v1.VariableV1(
+          1.0, dtype=dtypes.float32, trainable=False, name="var"
+      )
+
+      sess.run(variables.global_variables_initializer())
+      queue_runner_impl.start_queue_runners(sess)
+      sess.run(var.assign(3.0))
+
+  def testIsolateSessionState(self):
+    server = self._cached_server
+
+    with ops.Graph().as_default():
+      init_value = array_ops.placeholder(dtypes.int32)
+      v = variable_v1.VariableV1(init_value, validate_shape=False, name="v")
+
+      sharing_config = config_pb2.ConfigProto(isolate_session_state=False)
+      sharing_sess_0 = session.Session(server.target, config=sharing_config)
+      sharing_sess_1 = session.Session(server.target, config=sharing_config)
+
+      isolate_config = config_pb2.ConfigProto(isolate_session_state=True)
+      isolate_sess_0 = session.Session(server.target, config=isolate_config)
+      isolate_sess_1 = session.Session(server.target, config=isolate_config)
+
+      # Initially all variables are initialized.
+      for sess in [
+          sharing_sess_0,
+          sharing_sess_1,
+          isolate_sess_0,
+          isolate_sess_1,
+      ]:
+        with self.assertRaises(errors_impl.FailedPreconditionError):
+          sess.run(v)
+
+      # Shared sessions will see each other's updates, but isolated sessions
+      # will not.
+      sharing_sess_0.run(v.initializer, feed_dict={init_value: 86})
+      self.assertAllEqual(86, sharing_sess_0.run(v))
+      self.assertAllEqual(86, sharing_sess_1.run(v))
+      with self.assertRaises(errors_impl.FailedPreconditionError):
+        isolate_sess_0.run(v)
+      with self.assertRaises(errors_impl.FailedPreconditionError):
+        isolate_sess_1.run(v)
+
+      # Changing the shape works because `validate_shape` is False.
+      sharing_sess_1.run(v.initializer, feed_dict={init_value: [86, 99]})
+      self.assertAllEqual([86, 99], sharing_sess_0.run(v))
+      self.assertAllEqual([86, 99], sharing_sess_1.run(v))
+      with self.assertRaises(errors_impl.FailedPreconditionError):
+        isolate_sess_0.run(v)
+      with self.assertRaises(errors_impl.FailedPreconditionError):
+        isolate_sess_1.run(v)
+
+      # Initializing in an isolated session will only affect the state in that
+      # session.
+      isolate_sess_0.run(v.initializer, feed_dict={init_value: 37})
+      self.assertAllEqual([86, 99], sharing_sess_0.run(v))
+      self.assertAllEqual([86, 99], sharing_sess_1.run(v))
+      self.assertAllEqual(37, isolate_sess_0.run(v))
+      with self.assertRaises(errors_impl.FailedPreconditionError):
+        isolate_sess_1.run(v)
+
+      # Isolated sessions can have different shapes for the same variable.
+      isolate_sess_1.run(v.initializer, feed_dict={init_value: [19, 86]})
+      self.assertAllEqual([86, 99], sharing_sess_0.run(v))
+      self.assertAllEqual([86, 99], sharing_sess_1.run(v))
+      self.assertAllEqual(37, isolate_sess_0.run(v))
+      self.assertAllEqual([19, 86], isolate_sess_1.run(v))
+
+  def testShapeChangingIsolateState(self):
+    server = self._cached_server
+    sharing_config = config_pb2.ConfigProto(isolate_session_state=False)
+    isolate_config = config_pb2.ConfigProto(isolate_session_state=True)
+
+    with ops.Graph().as_default():
+      w_vector = variable_v1.VariableV1([1, 2, 3], name="w")
+      with session.Session(server.target, config=sharing_config) as sess:
+        with self.assertRaises(errors_impl.FailedPreconditionError):
+          sess.run(w_vector)
+        sess.run(w_vector.initializer)
+        self.assertAllEqual([1, 2, 3], sess.run(w_vector))
+
+    with ops.Graph().as_default():
+      w_vector = variable_v1.VariableV1([4, 5, 6], name="w")
+      with session.Session(server.target, config=sharing_config) as sess:
+        self.assertAllEqual([1, 2, 3], sess.run(w_vector))
+        sess.run(w_vector.initializer)
+        self.assertAllEqual([4, 5, 6], sess.run(w_vector))
+
+    with ops.Graph().as_default():
+      w_scalar = variable_v1.VariableV1(37, name="w")
+      with session.Session(server.target, config=isolate_config) as sess:
+        with self.assertRaises(errors_impl.FailedPreconditionError):
+          sess.run(w_scalar)
+        sess.run(w_scalar.initializer)
+        self.assertAllEqual(37, sess.run(w_scalar))
+
+
+class ServerDefTest(test.TestCase):
+
+  def testLocalServer(self):
+    cluster_def = server_lib.ClusterSpec(
+        {"local": ["localhost:2222"]}
+    ).as_cluster_def()
+    server_def = tensorflow_server_pb2.ServerDef(
+        cluster=cluster_def, job_name="local", task_index=0, protocol="grpc"
+    )
+
+    self.assertProtoEquals(
+        """
+    cluster {
+      job { name: 'local' tasks { key: 0 value: 'localhost:2222' } }
+    }
+    job_name: 'local' task_index: 0 protocol: 'grpc'
+    """,
+        server_def,
+    )
+
+    # Verifies round trip from Proto->Spec->Proto is correct.
+    cluster_spec = server_lib.ClusterSpec(cluster_def)
+    self.assertProtoEquals(cluster_def, cluster_spec.as_cluster_def())
+
+  def testTwoProcesses(self):
+    cluster_def = server_lib.ClusterSpec(
+        {"local": ["localhost:2222", "localhost:2223"]}
+    ).as_cluster_def()
+    server_def = tensorflow_server_pb2.ServerDef(
+        cluster=cluster_def, job_name="local", task_index=1, protocol="grpc"
+    )
+
+    self.assertProtoEquals(
+        """
+    cluster {
+      job { name: 'local' tasks { key: 0 value: 'localhost:2222' }
+                          tasks { key: 1 value: 'localhost:2223' } }
+    }
+    job_name: 'local' task_index: 1 protocol: 'grpc'
+    """,
+        server_def,
+    )
+
+    # Verifies round trip from Proto->Spec->Proto is correct.
+    cluster_spec = server_lib.ClusterSpec(cluster_def)
+    self.assertProtoEquals(cluster_def, cluster_spec.as_cluster_def())
+
+  def testTwoJobs(self):
+    cluster_def = server_lib.ClusterSpec({
+        "ps": ["ps0:2222", "ps1:2222"],
+        "worker": ["worker0:2222", "worker1:2222", "worker2:2222"],
+    }).as_cluster_def()
+    server_def = tensorflow_server_pb2.ServerDef(
+        cluster=cluster_def, job_name="worker", task_index=2, protocol="grpc"
+    )
+
+    self.assertProtoEquals(
+        """
+    cluster {
+      job { name: 'ps' tasks { key: 0 value: 'ps0:2222' }
+                       tasks { key: 1 value: 'ps1:2222' } }
+      job { name: 'worker' tasks { key: 0 value: 'worker0:2222' }
+                           tasks { key: 1 value: 'worker1:2222' }
+                           tasks { key: 2 value: 'worker2:2222' } }
+    }
+    job_name: 'worker' task_index: 2 protocol: 'grpc'
+    """,
+        server_def,
+    )
+
+    # Verifies round trip from Proto->Spec->Proto is correct.
+    cluster_spec = server_lib.ClusterSpec(cluster_def)
+    self.assertProtoEquals(cluster_def, cluster_spec.as_cluster_def())
+
+  def testDenseAndSparseJobs(self):
+    cluster_def = server_lib.ClusterSpec({
+        "ps": ["ps0:2222", "ps1:2222"],
+        "worker": {0: "worker0:2222", 2: "worker2:2222"},
+    }).as_cluster_def()
+    server_def = tensorflow_server_pb2.ServerDef(
+        cluster=cluster_def, job_name="worker", task_index=2, protocol="grpc"
+    )
+
+    self.assertProtoEquals(
+        """
+    cluster {
+      job { name: 'ps' tasks { key: 0 value: 'ps0:2222' }
+                       tasks { key: 1 value: 'ps1:2222' } }
+      job { name: 'worker' tasks { key: 0 value: 'worker0:2222' }
+                           tasks { key: 2 value: 'worker2:2222' } }
+    }
+    job_name: 'worker' task_index: 2 protocol: 'grpc'
+    """,
+        server_def,
+    )
+
+    # Verifies round trip from Proto->Spec->Proto is correct.
+    cluster_spec = server_lib.ClusterSpec(cluster_def)
+    self.assertProtoEquals(cluster_def, cluster_spec.as_cluster_def())
+
+
+class ClusterSpecTest(test.TestCase):
+
+  def testStringConversion(self):
+    cluster_spec = server_lib.ClusterSpec(
+        {"ps": ["ps0:1111"], "worker": ["worker0:3333", "worker1:4444"]}
+    )
+
+    expected_str = (
+        "ClusterSpec({'ps': ['ps0:1111'], 'worker': ['worker0:3333', "
+        "'worker1:4444']})"
+    )
+    self.assertEqual(expected_str, str(cluster_spec))
+
+  def testProtoDictDefEquivalences(self):
+    cluster_spec = server_lib.ClusterSpec({
+        "ps": ["ps0:2222", "ps1:2222"],
+        "worker": ["worker0:2222", "worker1:2222", "worker2:2222"],
+    })
+
+    expected_proto = """
+    job { name: 'ps' tasks { key: 0 value: 'ps0:2222' }
+                     tasks { key: 1 value: 'ps1:2222' } }
+    job { name: 'worker' tasks { key: 0 value: 'worker0:2222' }
+                         tasks { key: 1 value: 'worker1:2222' }
+                         tasks { key: 2 value: 'worker2:2222' } }
+    """
+
+    self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto, server_lib.ClusterSpec(cluster_spec).as_cluster_def()
+    )
+    self.assertProtoEquals(
+        expected_proto,
+        server_lib.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def(),
+    )
+    self.assertProtoEquals(
+        expected_proto,
+        server_lib.ClusterSpec(cluster_spec.as_dict()).as_cluster_def(),
+    )
+
+  def testProtoDictDefEquivalencesWithStringTaskIndex(self):
+    cluster_spec = server_lib.ClusterSpec(
+        {"ps": ["ps0:2222", "ps1:2222"], "worker": {"1": "worker1:2222"}}
+    )
+
+    expected_proto = """
+    job { name: 'ps' tasks { key: 0 value: 'ps0:2222' }
+                     tasks { key: 1 value: 'ps1:2222' } }
+    job { name: 'worker' tasks { key: 1 value: 'worker1:2222' } }
+    """
+
+    self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto, server_lib.ClusterSpec(cluster_spec).as_cluster_def()
+    )
+    self.assertProtoEquals(
+        expected_proto,
+        server_lib.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def(),
+    )
+    self.assertProtoEquals(
+        expected_proto,
+        server_lib.ClusterSpec(cluster_spec.as_dict()).as_cluster_def(),
+    )
+
+  def testProtoDictDefEquivalencesWithZeroWorker(self):
+    cluster_spec = server_lib.ClusterSpec(
+        {"ps": ["ps0:2222", "ps1:2222"], "worker": []}
+    )
+
+    expected_proto = """
+    job { name: 'ps' tasks { key: 0 value: 'ps0:2222' }
+                     tasks { key: 1 value: 'ps1:2222' } }
+    job { name: 'worker' }
+    """
+
+    self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto, server_lib.ClusterSpec(cluster_spec).as_cluster_def()
+    )
+    self.assertProtoEquals(
+        expected_proto,
+        server_lib.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def(),
+    )
+    self.assertProtoEquals(
+        expected_proto,
+        server_lib.ClusterSpec(cluster_spec.as_dict()).as_cluster_def(),
+    )
+
+  def testClusterSpecAccessors(self):
+    original_dict = {
+        "ps": ["ps0:2222", "ps1:2222"],
+        "worker": ["worker0:2222", "worker1:2222", "worker2:2222"],
+        "sparse": {0: "sparse0:2222", 3: "sparse3:2222"},
+    }
+    cluster_spec = server_lib.ClusterSpec(original_dict)
+
+    self.assertEqual(original_dict, cluster_spec.as_dict())
+
+    self.assertEqual(2, cluster_spec.num_tasks("ps"))
+    self.assertEqual(3, cluster_spec.num_tasks("worker"))
+    self.assertEqual(2, cluster_spec.num_tasks("sparse"))
+    with self.assertRaises(ValueError):
+      cluster_spec.num_tasks("unknown")
+
+    self.assertEqual("ps0:2222", cluster_spec.task_address("ps", 0))
+    self.assertEqual("sparse0:2222", cluster_spec.task_address("sparse", 0))
+    with self.assertRaises(ValueError):
+      cluster_spec.task_address("unknown", 0)
+    with self.assertRaises(ValueError):
+      cluster_spec.task_address("sparse", 2)
+
+    self.assertEqual([0, 1], cluster_spec.task_indices("ps"))
+    self.assertEqual([0, 1, 2], cluster_spec.task_indices("worker"))
+    self.assertEqual([0, 3], cluster_spec.task_indices("sparse"))
+    with self.assertRaises(ValueError):
+      cluster_spec.task_indices("unknown")
+
+    # NOTE(mrry): `ClusterSpec.job_tasks()` is not recommended for use
+    # with sparse jobs.
+    self.assertEqual(["ps0:2222", "ps1:2222"], cluster_spec.job_tasks("ps"))
+    self.assertEqual(
+        ["worker0:2222", "worker1:2222", "worker2:2222"],
+        cluster_spec.job_tasks("worker"),
+    )
+    self.assertEqual(
+        ["sparse0:2222", None, None, "sparse3:2222"],
+        cluster_spec.job_tasks("sparse"),
+    )
+    with self.assertRaises(ValueError):
+      cluster_spec.job_tasks("unknown")
+
+  def testEmptyClusterSpecIsFalse(self):
+    self.assertFalse(server_lib.ClusterSpec({}))
+
+  def testNonEmptyClusterSpecIsTrue(self):
+    self.assertTrue(server_lib.ClusterSpec({"job": ["host:port"]}))
+
+  def testEq(self):
+    self.assertEqual(server_lib.ClusterSpec({}), server_lib.ClusterSpec({}))
+    self.assertEqual(
+        server_lib.ClusterSpec({"job": ["host:2222"]}),
+        server_lib.ClusterSpec({"job": ["host:2222"]}),
+    )
+    self.assertEqual(
+        server_lib.ClusterSpec({"job": {0: "host:2222"}}),
+        server_lib.ClusterSpec({"job": ["host:2222"]}),
+    )
+
+  def testNe(self):
+    self.assertNotEqual(
+        server_lib.ClusterSpec({}),
+        server_lib.ClusterSpec({"job": ["host:2223"]}),
+    )
+    self.assertNotEqual(
+        server_lib.ClusterSpec({"job1": ["host:2222"]}),
+        server_lib.ClusterSpec({"job2": ["host:2222"]}),
+    )
+    self.assertNotEqual(
+        server_lib.ClusterSpec({"job": ["host:2222"]}),
+        server_lib.ClusterSpec({"job": ["host:2223"]}),
+    )
+    self.assertNotEqual(
+        server_lib.ClusterSpec({"job": ["host:2222", "host:2223"]}),
+        server_lib.ClusterSpec({"job": ["host:2223", "host:2222"]}),
+    )
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/session_manager.py

@@ -0,0 +1,607 @@
+# 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.
+# ==============================================================================
+"""Training helper that checkpoints models and creates session."""
+
+import time
+from typing import Optional, Tuple
+
+import numpy as np
+
+from tensorflow.python.checkpoint import checkpoint_management
+from tensorflow.python.client import session
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.training import saver as saver_lib
+from tensorflow.python.util.tf_export import tf_export
+
+
+def _maybe_name(obj) -> str:
+  """Returns object name if it has one, or a message otherwise.
+
+  This is useful for names that apper in error messages.
+  Args:
+    obj: Object to get the name of.
+  Returns:
+    name, "None", or a "no name" message.
+  """
+  if obj is None:
+    return "None"
+  elif hasattr(obj, "name"):
+    return obj.name
+  else:
+    return "<no name for %s>" % type(obj)
+
+
+def _restore_checkpoint_and_maybe_run_saved_model_initializers(
+    sess: session.Session, saver: saver_lib.Saver, path: str
+):
+  """Restores checkpoint values and SavedModel initializers if found."""
+  # NOTE: All references to SavedModel refer to SavedModels loaded from the
+  # load_v2 API (which does not require the `sess` argument).
+
+  # If the graph contains resources loaded from a SavedModel, they are not
+  # restored when calling `saver.restore`. Thus, the SavedModel initializer must
+  # be called with `saver.restore` to properly initialize the model.
+
+  # The SavedModel init is stored in the "saved_model_initializers" collection.
+  # This collection is part of the MetaGraph's default_init_op, so it is already
+  # called by MonitoredSession as long as the saver doesn't restore any
+  # checkpoints from the working dir.
+  saved_model_init_ops = ops.get_collection("saved_model_initializers")
+  if saved_model_init_ops:
+    sess.run(saved_model_init_ops)
+
+  # The saver must be called *after* the SavedModel init, because the SavedModel
+  # init will restore the variables from the SavedModel variables directory.
+  # Initializing/restoring twice is not ideal but there's no other way to do it.
+  saver.restore(sess, path)
+
+
+@tf_export(v1=["train.SessionManager"])
+class SessionManager:
+  """Training helper that restores from checkpoint and creates session.
+
+  This class is a small wrapper that takes care of session creation and
+  checkpoint recovery. It also provides functions that to facilitate
+  coordination among multiple training threads or processes.
+
+  * Checkpointing trained variables as the training progresses.
+  * Initializing variables on startup, restoring them from the most recent
+    checkpoint after a crash, or wait for checkpoints to become available.
+
+  ### Usage:
+
+  ```python
+  with tf.Graph().as_default():
+     ...add operations to the graph...
+    # Create a SessionManager that will checkpoint the model in '/tmp/mydir'.
+    sm = SessionManager()
+    sess = sm.prepare_session(master, init_op, saver, checkpoint_dir)
+    # Use the session to train the graph.
+    while True:
+      sess.run(<my_train_op>)
+  ```
+
+  `prepare_session()` initializes or restores a model. It requires `init_op`
+  and `saver` as an argument.
+
+  A second process could wait for the model to be ready by doing the following:
+
+  ```python
+  with tf.Graph().as_default():
+     ...add operations to the graph...
+    # Create a SessionManager that will wait for the model to become ready.
+    sm = SessionManager()
+    sess = sm.wait_for_session(master)
+    # Use the session to train the graph.
+    while True:
+      sess.run(<my_train_op>)
+  ```
+
+  `wait_for_session()` waits for a model to be initialized by other processes.
+
+  """
+
+  def __init__(
+      self,
+      local_init_op: ops.Operation = None,
+      ready_op: ops.Operation = None,
+      ready_for_local_init_op: ops.Operation = None,
+      graph: ops.Graph = None,
+      recovery_wait_secs=30,
+      local_init_run_options: "distribute_lib.RunOptions" = None,
+      local_init_feed_dict=None,
+  ):
+    """Creates a SessionManager.
+
+    The `local_init_op` is an `Operation` that is run always after a new session
+    was created. If `None`, this step is skipped.
+
+    The `ready_op` is an `Operation` used to check if the model is ready.  The
+    model is considered ready if that operation returns an empty 1D string
+    tensor. If the operation returns a non empty 1D string tensor, the elements
+    are concatenated and used to indicate to the user why the model is not
+    ready.
+
+    The `ready_for_local_init_op` is an `Operation` used to check if the model
+    is ready to run local_init_op.  The model is considered ready if that
+    operation returns an empty 1D string tensor. If the operation returns a non
+    empty 1D string tensor, the elements are concatenated and used to indicate
+    to the user why the model is not ready.
+
+    If `ready_op` is `None`, the model is not checked for readiness.
+
+    `recovery_wait_secs` is the number of seconds between checks that
+    the model is ready.  It is used by processes to wait for a model to
+    be initialized or restored.  Defaults to 30 seconds.
+
+    Args:
+      local_init_op: An `Operation` run immediately after session creation.
+         Usually used to initialize tables and local variables.
+      ready_op: An `Operation` to check if the model is initialized.
+      ready_for_local_init_op: An `Operation` to check if the model is ready
+         to run local_init_op.
+      graph: The `Graph` that the model will use.
+      recovery_wait_secs: Seconds between checks for the model to be ready.
+      local_init_run_options: RunOptions to be passed to session.run when
+        executing the local_init_op.
+      local_init_feed_dict: Optional session feed dictionary to use when running
+        the local_init_op.
+
+    Raises:
+      ValueError: If ready_for_local_init_op is not None but local_init_op is
+        None
+    """
+    # Sets default values of arguments.
+    if graph is None:
+      graph = ops.get_default_graph()
+    self._local_init_op = local_init_op
+    self._ready_op = ready_op
+    self._ready_for_local_init_op = ready_for_local_init_op
+    self._graph = graph
+    self._recovery_wait_secs = recovery_wait_secs
+    self._target = None
+    self._local_init_run_options = local_init_run_options
+    self._local_init_feed_dict = local_init_feed_dict
+    if ready_for_local_init_op is not None and local_init_op is None:
+      raise ValueError("If you pass a ready_for_local_init_op "
+                       "you must also pass a local_init_op "
+                       ", ready_for_local_init_op [%s]" %
+                       ready_for_local_init_op)
+
+  def _restore_checkpoint(
+      self,
+      master: str,
+      saver: saver_lib.Saver = None,
+      checkpoint_dir: str = None,
+      checkpoint_filename_with_path: str = None,
+      wait_for_checkpoint=False,
+      max_wait_secs=7200,
+      config=None,
+  ) -> Tuple[session.Session, bool]:
+    """Creates a `Session`, and tries to restore a checkpoint.
+
+
+    Args:
+      master: `String` representation of the TensorFlow master to use.
+      saver: A `Saver` object used to restore a model.
+      checkpoint_dir: Path to the checkpoint files. The latest checkpoint in the
+        dir will be used to restore.
+      checkpoint_filename_with_path: Full file name path to the checkpoint file.
+      wait_for_checkpoint: Whether to wait for checkpoint to become available.
+      max_wait_secs: Maximum time to wait for checkpoints to become available.
+      config: Optional `ConfigProto` proto used to configure the session.
+
+    Returns:
+      A pair (sess, is_restored) where 'is_restored' is `True` if
+      the session could be restored, `False` otherwise.
+
+    Raises:
+      ValueError: If both checkpoint_dir and checkpoint_filename_with_path are
+        set.
+    """
+    self._target = master
+
+    # This is required to so that we initialize the TPU device before
+    # restoring from checkpoint since we'll be placing variables on the device
+    # and TPUInitialize wipes out the memory of the device.
+    strategy = distribute_lib.get_strategy()
+    if strategy and hasattr(strategy.extended,
+                            "_experimental_initialize_system"):
+      strategy.extended._experimental_initialize_system()  # pylint: disable=protected-access
+
+    sess = session.Session(self._target, graph=self._graph, config=config)
+    if checkpoint_dir and checkpoint_filename_with_path:
+      raise ValueError("Can not provide both checkpoint_dir and "
+                       "checkpoint_filename_with_path.")
+    # If either saver or checkpoint_* is not specified, cannot restore. Just
+    # return.
+    if not saver or not (checkpoint_dir or checkpoint_filename_with_path):
+      return sess, False
+
+    if checkpoint_filename_with_path:
+      _restore_checkpoint_and_maybe_run_saved_model_initializers(
+          sess, saver, checkpoint_filename_with_path)
+      return sess, True
+
+    # Waits up until max_wait_secs for checkpoint to become available.
+    wait_time = 0
+    ckpt = checkpoint_management.get_checkpoint_state(checkpoint_dir)
+    while not ckpt or not ckpt.model_checkpoint_path:
+      if wait_for_checkpoint and wait_time < max_wait_secs:
+        logging.info("Waiting for checkpoint to be available.")
+        time.sleep(self._recovery_wait_secs)
+        wait_time += self._recovery_wait_secs
+        ckpt = checkpoint_management.get_checkpoint_state(checkpoint_dir)
+      else:
+        return sess, False
+
+    # Loads the checkpoint.
+    _restore_checkpoint_and_maybe_run_saved_model_initializers(
+        sess, saver, ckpt.model_checkpoint_path)
+    saver.recover_last_checkpoints(ckpt.all_model_checkpoint_paths)
+    return sess, True
+
+  def prepare_session(
+      self,
+      master: str,
+      init_op: ops.Operation = None,
+      saver: saver_lib.Saver = None,
+      checkpoint_dir: str = None,
+      checkpoint_filename_with_path: str = None,
+      wait_for_checkpoint=False,
+      max_wait_secs=7200,
+      config=None,
+      init_feed_dict=None,
+      init_fn=None,
+  ) -> session.Session:
+    """Creates a `Session`. Makes sure the model is ready to be used.
+
+    Creates a `Session` on 'master'. If a `saver` object is passed in, and
+    `checkpoint_dir` points to a directory containing valid checkpoint
+    files, then it will try to recover the model from checkpoint. If
+    no checkpoint files are available, and `wait_for_checkpoint` is
+    `True`, then the process would check every `recovery_wait_secs`,
+    up to `max_wait_secs`, for recovery to succeed.
+
+    If the model cannot be recovered successfully then it is initialized by
+    running the `init_op` and calling `init_fn` if they are provided.
+    The `local_init_op` is also run after init_op and init_fn, regardless of
+    whether the model was recovered successfully, but only if
+    `ready_for_local_init_op` passes.
+
+    If the model is recovered from a checkpoint it is assumed that all
+    global variables have been initialized, in particular neither `init_op`
+    nor `init_fn` will be executed.
+
+    It is an error if the model cannot be recovered and no `init_op`
+    or `init_fn` or `local_init_op` are passed.
+
+    Args:
+      master: `String` representation of the TensorFlow master to use.
+      init_op: Optional `Operation` used to initialize the model.
+      saver: A `Saver` object used to restore a model.
+      checkpoint_dir: Path to the checkpoint files. The latest checkpoint in the
+        dir will be used to restore.
+      checkpoint_filename_with_path: Full file name path to the checkpoint file.
+      wait_for_checkpoint: Whether to wait for checkpoint to become available.
+      max_wait_secs: Maximum time to wait for checkpoints to become available.
+      config: Optional `ConfigProto` proto used to configure the session.
+      init_feed_dict: Optional dictionary that maps `Tensor` objects to feed
+        values.  This feed dictionary is passed to the session `run()` call when
+        running the init op.
+      init_fn: Optional callable used to initialize the model. Called after the
+        optional `init_op` is called.  The callable must accept one argument,
+        the session being initialized.
+
+    Returns:
+      A `Session` object that can be used to drive the model.
+
+    Raises:
+      RuntimeError: If the model cannot be initialized or recovered.
+      ValueError: If both checkpoint_dir and checkpoint_filename_with_path are
+        set.
+    """
+
+    sess, is_loaded_from_checkpoint = self._restore_checkpoint(
+        master,
+        saver,
+        checkpoint_dir=checkpoint_dir,
+        checkpoint_filename_with_path=checkpoint_filename_with_path,
+        wait_for_checkpoint=wait_for_checkpoint,
+        max_wait_secs=max_wait_secs,
+        config=config)
+    if not is_loaded_from_checkpoint:
+      if init_op is None and not init_fn and self._local_init_op is None:
+        raise RuntimeError("Model is not initialized and no init_op or "
+                           "init_fn or local_init_op was given")
+      if init_op is not None:
+        sess.run(init_op, feed_dict=init_feed_dict)
+      if init_fn:
+        init_fn(sess)
+
+    local_init_success, msg = self._try_run_local_init_op(sess)
+    if not local_init_success:
+      raise RuntimeError(
+          "Init operations did not make model ready for local_init.  "
+          "Init op: %s, init fn: %s, error: %s" % (_maybe_name(init_op),
+                                                   init_fn,
+                                                   msg))
+
+    is_ready, msg = self._model_ready(sess)
+    if not is_ready:
+      raise RuntimeError(
+          "Init operations did not make model ready.  "
+          "Init op: %s, init fn: %s, local_init_op: %s, error: %s" %
+          (_maybe_name(init_op), init_fn, self._local_init_op, msg))
+    return sess
+
+  def recover_session(
+      self,
+      master: str,
+      saver: saver_lib.Saver = None,
+      checkpoint_dir: str = None,
+      checkpoint_filename_with_path: str = None,
+      wait_for_checkpoint=False,
+      max_wait_secs=7200,
+      config=None,
+  ) -> Tuple[session.Session, bool]:
+    """Creates a `Session`, recovering if possible.
+
+    Creates a new session on 'master'.  If the session is not initialized
+    and can be recovered from a checkpoint, recover it.
+
+    Args:
+      master: `String` representation of the TensorFlow master to use.
+      saver: A `Saver` object used to restore a model.
+      checkpoint_dir: Path to the checkpoint files. The latest checkpoint in the
+        dir will be used to restore.
+      checkpoint_filename_with_path: Full file name path to the checkpoint file.
+      wait_for_checkpoint: Whether to wait for checkpoint to become available.
+      max_wait_secs: Maximum time to wait for checkpoints to become available.
+      config: Optional `ConfigProto` proto used to configure the session.
+
+    Returns:
+      A pair (sess, initialized) where 'initialized' is `True` if
+      the session could be recovered and initialized, `False` otherwise.
+
+    Raises:
+      ValueError: If both checkpoint_dir and checkpoint_filename_with_path are
+        set.
+    """
+
+    sess, is_loaded_from_checkpoint = self._restore_checkpoint(
+        master,
+        saver,
+        checkpoint_dir=checkpoint_dir,
+        checkpoint_filename_with_path=checkpoint_filename_with_path,
+        wait_for_checkpoint=wait_for_checkpoint,
+        max_wait_secs=max_wait_secs,
+        config=config)
+
+    # Always try to run local_init_op
+    local_init_success, msg = self._try_run_local_init_op(sess)
+
+    if not is_loaded_from_checkpoint:
+      # Do not need to run checks for readiness
+      return sess, False
+
+    restoring_file = checkpoint_dir or checkpoint_filename_with_path
+    if not local_init_success:
+      logging.info(
+          "Restoring model from %s did not make model ready for local init:"
+          " %s", restoring_file, msg)
+      return sess, False
+
+    is_ready, msg = self._model_ready(sess)
+    if not is_ready:
+      logging.info("Restoring model from %s did not make model ready: %s",
+                   restoring_file, msg)
+      return sess, False
+
+    logging.info("Restored model from %s", restoring_file)
+    return sess, is_loaded_from_checkpoint
+
+  def wait_for_session(
+      self, master: str, config=None, max_wait_secs=float("Inf")
+  ) -> Optional[session.Session]:
+    """Creates a new `Session` and waits for model to be ready.
+
+    Creates a new `Session` on 'master'.  Waits for the model to be
+    initialized or recovered from a checkpoint.  It's expected that
+    another thread or process will make the model ready, and that this
+    is intended to be used by threads/processes that participate in a
+    distributed training configuration where a different thread/process
+    is responsible for initializing or recovering the model being trained.
+
+    NB: The amount of time this method waits for the session is bounded
+    by max_wait_secs. By default, this function will wait indefinitely.
+
+    Args:
+      master: `String` representation of the TensorFlow master to use.
+      config: Optional ConfigProto proto used to configure the session.
+      max_wait_secs: Maximum time to wait for the session to become available.
+
+    Returns:
+      A `Session`. May be None if the operation exceeds the timeout
+      specified by config.operation_timeout_in_ms.
+
+    Raises:
+      tf.DeadlineExceededError: if the session is not available after
+        max_wait_secs.
+    """
+    self._target = master
+
+    if max_wait_secs is None:
+      max_wait_secs = float("Inf")
+    timer = _CountDownTimer(max_wait_secs)
+
+    while True:
+      sess = session.Session(self._target, graph=self._graph, config=config)
+      not_ready_msg = None
+      not_ready_local_msg = None
+      local_init_success, not_ready_local_msg = self._try_run_local_init_op(
+          sess)
+      if local_init_success:
+        # Successful if local_init_op is None, or ready_for_local_init_op passes
+        is_ready, not_ready_msg = self._model_ready(sess)
+        if is_ready:
+          return sess
+
+      self._safe_close(sess)
+
+      # Do we have enough time left to try again?
+      remaining_ms_after_wait = (
+          timer.secs_remaining() - self._recovery_wait_secs)
+      if remaining_ms_after_wait < 0:
+        raise errors.DeadlineExceededError(
+            None, None,
+            "Session was not ready after waiting %d secs." % (max_wait_secs,))
+
+      logging.info("Waiting for model to be ready.  "
+                   "Ready_for_local_init_op:  %s, ready: %s",
+                   not_ready_local_msg, not_ready_msg)
+      time.sleep(self._recovery_wait_secs)
+
+  def _safe_close(self, sess: session.Session):
+    """Closes a session without raising an exception.
+
+    Just like sess.close() but ignores exceptions.
+
+    Args:
+      sess: A `Session`.
+    """
+    # pylint: disable=broad-except
+    try:
+      sess.close()
+    except Exception:
+      # Intentionally not logging to avoid user complaints that
+      # they get cryptic errors.  We really do not care that Close
+      # fails.
+      pass
+    # pylint: enable=broad-except
+
+  def _model_ready(self, sess: session.Session) -> Tuple[bool, Optional[str]]:
+    """Checks if the model is ready or not.
+
+    Args:
+      sess: A `Session`.
+
+    Returns:
+      A tuple (is_ready, msg), where is_ready is True if ready and False
+      otherwise, and msg is `None` if the model is ready, a `String` with the
+      reason why it is not ready otherwise.
+    """
+    return _ready(self._ready_op, sess, "Model not ready")
+
+  def _model_ready_for_local_init(
+      self, sess: session.Session
+  ) -> Tuple[bool, Optional[str]]:
+    """Checks if the model is ready to run local_init_op.
+
+    Args:
+      sess: A `Session`.
+
+    Returns:
+      A tuple (is_ready, msg), where is_ready is True if ready to run
+      local_init_op and False otherwise, and msg is `None` if the model is
+      ready to run local_init_op, a `String` with the reason why it is not ready
+      otherwise.
+    """
+    return _ready(self._ready_for_local_init_op, sess,
+                  "Model not ready for local init")
+
+  def _try_run_local_init_op(
+      self, sess: session.Session
+  ) -> Tuple[bool, Optional[str]]:
+    """Tries to run _local_init_op, if not None, and is ready for local init.
+
+    Args:
+      sess: A `Session`.
+
+    Returns:
+      A tuple (is_successful, msg), where is_successful is True if
+      _local_init_op is None, or we ran _local_init_op, and False otherwise;
+      and msg is a `String` with the reason why the model was not ready to run
+      local init.
+    """
+    if self._local_init_op is not None:
+      is_ready_for_local_init, msg = self._model_ready_for_local_init(sess)
+      if is_ready_for_local_init:
+        logging.info("Running local_init_op.")
+        sess.run(self._local_init_op, feed_dict=self._local_init_feed_dict,
+                 options=self._local_init_run_options)
+        logging.info("Done running local_init_op.")
+        return True, None
+      else:
+        return False, msg
+    return True, None
+
+
+def _ready(
+    op: ops.Operation, sess: session.Session, msg
+) -> Tuple[bool, Optional[str]]:
+  """Checks if the model is ready or not, as determined by op.
+
+  Args:
+    op: An op, either _ready_op or _ready_for_local_init_op, which defines the
+      readiness of the model.
+    sess: A `Session`.
+    msg: A message to log to warning if not ready
+
+  Returns:
+    A tuple (is_ready, msg), where is_ready is True if ready and False
+    otherwise, and msg is `None` if the model is ready, a `String` with the
+    reason why it is not ready otherwise.
+  """
+  if op is None:
+    return True, None
+  else:
+    try:
+      ready_value = sess.run(op)
+      # The model is considered ready if ready_op returns an empty 1-D tensor.
+      # Also compare to `None` and dtype being int32 for backward
+      # compatibility.
+      if (ready_value is None or ready_value.dtype == np.int32 or
+          ready_value.size == 0):
+        return True, None
+      else:
+        # TODO(sherrym): If a custom ready_op returns other types of tensor,
+        # or strings other than variable names, this message could be
+        # confusing.
+        non_initialized_varnames = ", ".join(
+            [i.decode("utf-8") for i in ready_value])
+        return False, "Variables not initialized: " + non_initialized_varnames
+    except errors.FailedPreconditionError as e:
+      if "uninitialized" not in str(e):
+        logging.warning("%s : error [%s]", msg, str(e))
+        raise e
+      return False, str(e)
+
+
+class _CountDownTimer:
+  """A timer that tracks a duration since creation."""
+
+  __slots__ = ["_start_time_secs", "_duration_secs"]
+
+  def __init__(self, duration_secs):
+    self._start_time_secs = time.time()
+    self._duration_secs = duration_secs
+
+  def secs_remaining(self):
+    diff = self._duration_secs - (time.time() - self._start_time_secs)
+    return max(0, diff)

tensorflow/python/training/session_manager_test.py

@@ -0,0 +1,825 @@
+# 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.
+# ==============================================================================
+"""Tests for SessionManager."""
+
+import os
+from tensorflow.python.checkpoint import checkpoint_management
+from tensorflow.python.client import session as session_lib
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import resource_variables_toggle
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.ops import variables
+from tensorflow.python.ops import while_loop
+from tensorflow.python.platform import gfile
+from tensorflow.python.platform import test
+from tensorflow.python.training import saver as saver_lib
+from tensorflow.python.training import server_lib
+from tensorflow.python.training import session_manager
+
+
+class SessionManagerTest(test.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    super(SessionManagerTest, cls).setUpClass()
+    resource_variables_toggle.disable_resource_variables()
+
+  def testPrepareSessionSucceeds(self):
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1([1.0, 2.0, 3.0], name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables())
+      sess = sm.prepare_session(
+          "", init_op=variables.global_variables_initializer())
+      self.assertAllClose([1.0, 2.0, 3.0], sess.run(v))
+
+  def testPrepareSessionSucceedsWithInitFeedDict(self):
+    with ops.Graph().as_default():
+      p = array_ops.placeholder(dtypes.float32, shape=(3,))
+      v = variable_v1.VariableV1(p, name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables())
+      sess = sm.prepare_session(
+          "",
+          init_op=variables.global_variables_initializer(),
+          init_feed_dict={p: [1.0, 2.0, 3.0]})
+      self.assertAllClose([1.0, 2.0, 3.0], sess.run(v))
+
+  def testPrepareSessionSucceedsWithInitFn(self):
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1([125], name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables())
+      sess = sm.prepare_session(
+          "", init_fn=lambda sess: sess.run(v.initializer))
+      self.assertAllClose([125], sess.run(v))
+
+  def testPrepareSessionSucceedsWithLocalInitFeedDict(self):
+    with ops.Graph().as_default():
+      p = array_ops.placeholder(dtypes.float32, shape=(3,))
+      v = variable_v1.VariableV1(
+          p, name="v", collections=[ops.GraphKeys.LOCAL_VARIABLES])
+      sm = session_manager.SessionManager(
+          local_init_op=v.initializer,
+          local_init_feed_dict={p: [1.0, 2.0, 3.0]},
+          ready_op=variables.report_uninitialized_variables())
+      sess = sm.prepare_session("")
+      self.assertAllClose([1.0, 2.0, 3.0], sess.run(v))
+
+  def testPrepareSessionFails(self):
+    checkpoint_dir = os.path.join(self.get_temp_dir(), "prepare_session")
+    checkpoint_dir2 = os.path.join(self.get_temp_dir(), "prepare_session2")
+    try:
+      gfile.DeleteRecursively(checkpoint_dir)
+      gfile.DeleteRecursively(checkpoint_dir2)
+    except errors.OpError:
+      pass  # Ignore
+    gfile.MakeDirs(checkpoint_dir)
+
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1([1.0, 2.0, 3.0], name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables())
+      saver = saver_lib.Saver({"v": v})
+      sess = sm.prepare_session(
+          "",
+          init_op=variables.global_variables_initializer(),
+          saver=saver,
+          checkpoint_dir=checkpoint_dir)
+      self.assertAllClose([1.0, 2.0, 3.0], sess.run(v))
+      checkpoint_filename = os.path.join(checkpoint_dir,
+                                         "prepare_session_checkpoint")
+      saver.save(sess, checkpoint_filename)
+    # Create a new Graph and SessionManager and recover.
+    with ops.Graph().as_default():
+      # Renames the checkpoint directory.
+      os.rename(checkpoint_dir, checkpoint_dir2)
+      gfile.MakeDirs(checkpoint_dir)
+      v = variable_v1.VariableV1([6.0, 7.0, 8.0], name="v")
+      with self.cached_session():
+        self.assertEqual(False, variable_v1.is_variable_initialized(v).eval())
+      session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables())
+      saver = saver_lib.Saver({"v": v})
+      # This should fail as there's no checkpoint within 2 seconds.
+      with self.assertRaisesRegex(
+          RuntimeError, "no init_op or init_fn or local_init_op was given"):
+        sess = sm.prepare_session(
+            "",
+            init_op=None,
+            saver=saver,
+            checkpoint_dir=checkpoint_dir,
+            wait_for_checkpoint=True,
+            max_wait_secs=2)
+      # Rename the checkpoint directory back.
+      gfile.DeleteRecursively(checkpoint_dir)
+      os.rename(checkpoint_dir2, checkpoint_dir)
+      # This should succeed as there's checkpoint.
+      sess = sm.prepare_session(
+          "",
+          init_op=None,
+          saver=saver,
+          checkpoint_dir=checkpoint_dir,
+          wait_for_checkpoint=True,
+          max_wait_secs=2)
+      self.assertEqual(
+          True,
+          variable_v1.is_variable_initialized(
+              sess.graph.get_tensor_by_name("v:0")).eval(session=sess))
+
+  def _test_recovered_variable(self,
+                               checkpoint_dir=None,
+                               checkpoint_filename_with_path=None):
+    # Create a new Graph and SessionManager and recover from a checkpoint.
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1(2, name="v")
+      with session_lib.Session():
+        self.assertEqual(False, variable_v1.is_variable_initialized(v).eval())
+      sm2 = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables())
+      saver = saver_lib.Saver({"v": v})
+      sess, initialized = sm2.recover_session(
+          "",
+          saver=saver,
+          checkpoint_dir=checkpoint_dir,
+          checkpoint_filename_with_path=checkpoint_filename_with_path)
+      self.assertTrue(initialized)
+      self.assertEqual(
+          True,
+          variable_v1.is_variable_initialized(
+              sess.graph.get_tensor_by_name("v:0")).eval(session=sess))
+      self.assertEqual(1, sess.run(v))
+
+  def testRecoverSession(self):
+    # Create a checkpoint.
+    checkpoint_dir = os.path.join(self.get_temp_dir(), "recover_session")
+    try:
+      gfile.DeleteRecursively(checkpoint_dir)
+    except errors.OpError:
+      pass  # Ignore
+    gfile.MakeDirs(checkpoint_dir)
+
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1(1, name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables())
+      saver = saver_lib.Saver({"v": v})
+      sess, initialized = sm.recover_session(
+          "", saver=saver, checkpoint_dir=checkpoint_dir)
+      self.assertFalse(initialized)
+      sess.run(v.initializer)
+      self.assertEqual(1, sess.run(v))
+      saver.save(sess, os.path.join(checkpoint_dir,
+                                    "recover_session_checkpoint"))
+    self._test_recovered_variable(checkpoint_dir=checkpoint_dir)
+    self._test_recovered_variable(
+        checkpoint_filename_with_path=checkpoint_management.latest_checkpoint(
+            checkpoint_dir))
+    # Cannot set both checkpoint_dir and checkpoint_filename_with_path.
+    with self.assertRaises(ValueError):
+      self._test_recovered_variable(
+          checkpoint_dir=checkpoint_dir,
+          checkpoint_filename_with_path=checkpoint_management.latest_checkpoint(
+              checkpoint_dir))
+
+  def testWaitForSessionReturnsNoneAfterTimeout(self):
+    with ops.Graph().as_default():
+      variable_v1.VariableV1(1, name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables(),
+          recovery_wait_secs=1)
+
+      # Set max_wait_secs to allow us to try a few times.
+      with self.assertRaises(errors.DeadlineExceededError):
+        sm.wait_for_session(master="", max_wait_secs=3)
+
+  def testInitWithNoneLocalInitOpError(self):
+    # Creating a SessionManager with a None local_init_op but
+    # non-None ready_for_local_init_op raises ValueError
+    with self.assertRaisesRegex(
+        ValueError, "If you pass a ready_for_local_init_op "
+        "you must also pass a local_init_op "):
+      session_manager.SessionManager(
+          ready_for_local_init_op=variables.report_uninitialized_variables(
+              variables.global_variables()),
+          local_init_op=None)
+
+  def testRecoverSessionWithReadyForLocalInitOp(self):
+    # Create a checkpoint.
+    checkpoint_dir = os.path.join(self.get_temp_dir(),
+                                  "recover_session_ready_for_local_init")
+    try:
+      gfile.DeleteRecursively(checkpoint_dir)
+    except errors.OpError:
+      pass  # Ignore
+    gfile.MakeDirs(checkpoint_dir)
+
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1(1, name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables())
+      saver = saver_lib.Saver({"v": v})
+      sess, initialized = sm.recover_session(
+          "", saver=saver, checkpoint_dir=checkpoint_dir)
+      self.assertFalse(initialized)
+      sess.run(v.initializer)
+      self.assertEqual(1, sess.run(v))
+      saver.save(sess, os.path.join(checkpoint_dir,
+                                    "recover_session_checkpoint"))
+    # Create a new Graph and SessionManager and recover.
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1(2, name="v")
+      w = variable_v1.VariableV1(
+          v,
+          trainable=False,
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
+          name="w")
+      with self.cached_session():
+        self.assertEqual(False, variable_v1.is_variable_initialized(v).eval())
+        self.assertEqual(False, variable_v1.is_variable_initialized(w).eval())
+      sm2 = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables(),
+          ready_for_local_init_op=variables.report_uninitialized_variables(
+              variables.global_variables()),
+          local_init_op=w.initializer)
+      saver = saver_lib.Saver({"v": v})
+      sess, initialized = sm2.recover_session(
+          "", saver=saver, checkpoint_dir=checkpoint_dir)
+      self.assertTrue(initialized)
+      self.assertEqual(
+          True,
+          variable_v1.is_variable_initialized(
+              sess.graph.get_tensor_by_name("v:0")).eval(session=sess))
+      self.assertEqual(
+          True,
+          variable_v1.is_variable_initialized(
+              sess.graph.get_tensor_by_name("w:0")).eval(session=sess))
+      self.assertEqual(1, sess.run(v))
+      self.assertEqual(1, sess.run(w))
+
+  def testRecoverSessionWithReadyForLocalInitOpFailsToReadyLocal(self):
+    # We use ready_for_local_init_op=report_uninitialized_variables(),
+    # which causes recover_session to not run local_init_op, and to return
+    # initialized=False
+
+    # Create a checkpoint.
+    checkpoint_dir = os.path.join(
+        self.get_temp_dir(),
+        "recover_session_ready_for_local_init_fails_to_ready_local")
+    try:
+      gfile.DeleteRecursively(checkpoint_dir)
+    except errors.OpError:
+      pass  # Ignore
+    gfile.MakeDirs(checkpoint_dir)
+
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1(1, name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables())
+      saver = saver_lib.Saver({"v": v})
+      sess, initialized = sm.recover_session(
+          "", saver=saver, checkpoint_dir=checkpoint_dir)
+      self.assertFalse(initialized)
+      sess.run(v.initializer)
+      self.assertEqual(1, sess.run(v))
+      saver.save(sess, os.path.join(checkpoint_dir,
+                                    "recover_session_checkpoint"))
+    # Create a new Graph and SessionManager and recover.
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1(2, name="v")
+      w = variable_v1.VariableV1(
+          v,
+          trainable=False,
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
+          name="w")
+      with self.cached_session():
+        self.assertEqual(False, variable_v1.is_variable_initialized(v).eval())
+        self.assertEqual(False, variable_v1.is_variable_initialized(w).eval())
+      sm2 = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables(),
+          ready_for_local_init_op=variables.report_uninitialized_variables(),
+          local_init_op=w.initializer)
+      saver = saver_lib.Saver({"v": v})
+      sess, initialized = sm2.recover_session(
+          "", saver=saver, checkpoint_dir=checkpoint_dir)
+      self.assertFalse(initialized)
+      self.assertEqual(
+          True,
+          variable_v1.is_variable_initialized(
+              sess.graph.get_tensor_by_name("v:0")).eval(session=sess))
+      self.assertEqual(
+          False,
+          variable_v1.is_variable_initialized(
+              sess.graph.get_tensor_by_name("w:0")).eval(session=sess))
+      self.assertEqual(1, sess.run(v))
+
+  def testRecoverSessionNoChkptStillRunsLocalInitOp(self):
+    # This test checks for backwards compatibility.
+    # In particular, we continue to ensure that recover_session will execute
+    # local_init_op exactly once, regardless of whether the session was
+    # successfully recovered.
+    with ops.Graph().as_default():
+      w = variable_v1.VariableV1(
+          1,
+          trainable=False,
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
+          name="w")
+      with self.cached_session():
+        self.assertEqual(False, variable_v1.is_variable_initialized(w).eval())
+      sm2 = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables(),
+          ready_for_local_init_op=None,
+          local_init_op=w.initializer)
+      # Try to recover session from None
+      sess, initialized = sm2.recover_session(
+          "", saver=None, checkpoint_dir=None)
+      # Succeeds because recover_session still run local_init_op
+      self.assertFalse(initialized)
+      self.assertEqual(
+          True,
+          variable_v1.is_variable_initialized(
+              sess.graph.get_tensor_by_name("w:0")).eval(session=sess))
+      self.assertEqual(1, sess.run(w))
+
+  def testRecoverSessionFailsStillRunsLocalInitOp(self):
+    # Create a checkpoint.
+    checkpoint_dir = os.path.join(
+        self.get_temp_dir(),
+        "recover_session_ready_for_local_init_fails_stil_run")
+    try:
+      gfile.DeleteRecursively(checkpoint_dir)
+    except errors.OpError:
+      pass  # Ignore
+    gfile.MakeDirs(checkpoint_dir)
+
+    # Create a new Graph and SessionManager and recover.
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1(2, name="v")
+      w = variable_v1.VariableV1(
+          1,
+          trainable=False,
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
+          name="w")
+      with self.cached_session():
+        self.assertEqual(False, variable_v1.is_variable_initialized(v).eval())
+        self.assertEqual(False, variable_v1.is_variable_initialized(w).eval())
+      sm2 = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables(),
+          ready_for_local_init_op=None,
+          local_init_op=w.initializer)
+      saver = saver_lib.Saver({"v": v})
+      sess, initialized = sm2.recover_session(
+          "",
+          saver=saver,
+          checkpoint_dir=checkpoint_dir,
+          wait_for_checkpoint=False)
+      self.assertFalse(initialized)
+      self.assertEqual(
+          False,
+          variable_v1.is_variable_initialized(
+              sess.graph.get_tensor_by_name("v:0")).eval(session=sess))
+      self.assertEqual(
+          True,
+          variable_v1.is_variable_initialized(
+              sess.graph.get_tensor_by_name("w:0")).eval(session=sess))
+      self.assertEqual(1, sess.run(w))
+
+  def testWaitForSessionLocalInit(self):
+    server = server_lib.Server.create_local_server()
+    with ops.Graph().as_default() as graph:
+      v = variable_v1.VariableV1(1, name="v")
+      w = variable_v1.VariableV1(
+          v,
+          trainable=False,
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
+          name="w")
+      sm = session_manager.SessionManager(
+          graph=graph,
+          ready_op=variables.report_uninitialized_variables(),
+          ready_for_local_init_op=variables.report_uninitialized_variables(
+              variables.global_variables()),
+          local_init_op=w.initializer)
+
+      # Initialize v but not w
+      s = session_lib.Session(server.target, graph=graph)
+      s.run(v.initializer)
+
+      sess = sm.wait_for_session(server.target, max_wait_secs=3)
+      self.assertEqual(
+          True,
+          variable_v1.is_variable_initialized(
+              sess.graph.get_tensor_by_name("v:0")).eval(session=sess))
+      self.assertEqual(
+          True,
+          variable_v1.is_variable_initialized(
+              sess.graph.get_tensor_by_name("w:0")).eval(session=sess))
+      self.assertEqual(1, sess.run(v))
+      self.assertEqual(1, sess.run(w))
+
+  def testWaitForSessionWithReadyForLocalInitOpFailsToReadyLocal(self):
+    with ops.Graph().as_default() as graph:
+      v = variable_v1.VariableV1(1, name="v")
+      w = variable_v1.VariableV1(
+          v,
+          trainable=False,
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
+          name="w")
+      sm = session_manager.SessionManager(
+          graph=graph,
+          ready_op=variables.report_uninitialized_variables(),
+          ready_for_local_init_op=variables.report_uninitialized_variables(),
+          local_init_op=w.initializer)
+
+      with self.assertRaises(errors_impl.DeadlineExceededError):
+        # Time-out because w fails to be initialized,
+        # because of overly restrictive ready_for_local_init_op
+        sm.wait_for_session("", max_wait_secs=3)
+
+  @test_util.run_v1_only("Requires TF V1 variable behavior.")
+  def testWaitForSessionInsufficientReadyForLocalInitCheck(self):
+    with ops.Graph().as_default() as graph:
+      v = variable_v1.VariableV1(1, name="v")
+      w = variable_v1.VariableV1(
+          v,
+          trainable=False,
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
+          name="w")
+      sm = session_manager.SessionManager(
+          graph=graph,
+          ready_op=variables.report_uninitialized_variables(),
+          ready_for_local_init_op=None,
+          local_init_op=w.initializer)
+    with self.assertRaisesRegex(errors_impl.DeadlineExceededError,
+                                "Session was not ready after waiting.*"):
+      sm.wait_for_session("", max_wait_secs=3)
+
+  def testPrepareSessionWithReadyForLocalInitOp(self):
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1(1, name="v")
+      w = variable_v1.VariableV1(
+          v,
+          trainable=False,
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
+          name="w")
+      x = variable_v1.VariableV1(
+          3 * v,
+          trainable=False,
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
+          name="x")
+      with self.cached_session():
+        self.assertEqual(False, variable_v1.is_variable_initialized(v).eval())
+        self.assertEqual(False, variable_v1.is_variable_initialized(w).eval())
+        self.assertEqual(False, variable_v1.is_variable_initialized(x).eval())
+      sm2 = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables(),
+          ready_for_local_init_op=variables.report_uninitialized_variables(
+              variables.global_variables()),
+          local_init_op=[w.initializer, x.initializer])
+      sess = sm2.prepare_session("", init_op=v.initializer)
+      self.assertEqual(
+          True,
+          variable_v1.is_variable_initialized(
+              sess.graph.get_tensor_by_name("v:0")).eval(session=sess))
+      self.assertEqual(
+          True,
+          variable_v1.is_variable_initialized(
+              sess.graph.get_tensor_by_name("w:0")).eval(session=sess))
+      self.assertEqual(
+          True,
+          variable_v1.is_variable_initialized(
+              sess.graph.get_tensor_by_name("x:0")).eval(session=sess))
+      self.assertEqual(1, sess.run(v))
+      self.assertEqual(1, sess.run(w))
+      self.assertEqual(3, sess.run(x))
+
+  @test_util.run_v1_only("Requires TF V1 variable behavior.")
+  def testPrepareSessionWithPartialInitOp(self):
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1(1, name="v")
+      w = variable_v1.VariableV1(
+          v,
+          trainable=False,
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
+          name="w")
+      x = variable_v1.VariableV1(
+          3 * v,
+          trainable=False,
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
+          name="x")
+      # TODO(b/70206927): Use ResourceVariables once they are handled properly.
+      v_res = variable_v1.VariableV1(1, name="v_res")
+      w_res = variable_v1.VariableV1(
+          v_res,
+          trainable=False,
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
+          name="w_res")
+      x_res = variable_v1.VariableV1(
+          3 * v_res,
+          trainable=False,
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
+          name="x_res")
+
+      with self.cached_session():
+        self.assertEqual(False, variable_v1.is_variable_initialized(v).eval())
+        self.assertEqual(False, variable_v1.is_variable_initialized(w).eval())
+        self.assertEqual(False, variable_v1.is_variable_initialized(x).eval())
+        self.assertEqual(False,
+                         variable_v1.is_variable_initialized(v_res).eval())
+        self.assertEqual(False,
+                         variable_v1.is_variable_initialized(w_res).eval())
+        self.assertEqual(False,
+                         variable_v1.is_variable_initialized(x_res).eval())
+      sm2 = session_manager.SessionManager(local_init_op=[
+          w.initializer, x.initializer, w_res.initializer, x_res.initializer
+      ])
+      sess = sm2.prepare_session("", init_op=None)
+      self.assertEqual(
+          False,
+          variable_v1.is_variable_initialized(
+              sess.graph.get_tensor_by_name("v:0")).eval(session=sess))
+      self.assertEqual(
+          True,
+          variable_v1.is_variable_initialized(
+              sess.graph.get_tensor_by_name("w:0")).eval(session=sess))
+      self.assertEqual(
+          True,
+          variable_v1.is_variable_initialized(
+              sess.graph.get_tensor_by_name("x:0")).eval(session=sess))
+      self.assertEqual(1, sess.run(w))
+      self.assertEqual(3, sess.run(x))
+      self.assertEqual(
+          False,
+          variable_v1.is_variable_initialized(
+              sess.graph.get_tensor_by_name("v_res:0")).eval(session=sess))
+      self.assertEqual(
+          True,
+          variable_v1.is_variable_initialized(
+              sess.graph.get_tensor_by_name("w_res:0")).eval(session=sess))
+      self.assertEqual(
+          True,
+          variable_v1.is_variable_initialized(
+              sess.graph.get_tensor_by_name("x_res:0")).eval(session=sess))
+      self.assertEqual(1, sess.run(w_res))
+      self.assertEqual(3, sess.run(x_res))
+
+  def testPrepareSessionWithCyclicInitializer(self):
+    # Regression test. Previously Variable._build_initializer_expr would enter
+    # into an infinite recursion when the variable's initial_value involved
+    # cyclic dependencies.
+    with ops.Graph().as_default():
+      i = while_loop.while_loop(lambda i: i < 1, lambda i: i + 1, [0])
+      v = variable_v1.VariableV1(array_ops.identity(i), name="v")
+      with self.cached_session():
+        self.assertEqual(False, variable_v1.is_variable_initialized(v).eval())
+      sm = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables())
+      sess = sm.prepare_session("", init_op=v.initializer)
+      self.assertEqual(1, sess.run(v))
+      self.assertEqual(
+          True,
+          variable_v1.is_variable_initialized(
+              sess.graph.get_tensor_by_name("v:0")).eval(session=sess))
+
+  def testPrepareSessionDidNotInitLocalVariable(self):
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1(1, name="v")
+      w = variable_v1.VariableV1(
+          v,
+          trainable=False,
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
+          name="w")
+      with self.cached_session():
+        self.assertEqual(False, variable_v1.is_variable_initialized(v).eval())
+        self.assertEqual(False, variable_v1.is_variable_initialized(w).eval())
+      sm2 = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables())
+      with self.assertRaisesRegex(RuntimeError,
+                                  "Init operations did not make model ready.*"):
+        sm2.prepare_session("", init_op=v.initializer)
+
+  def testPrepareSessionDidNotInitLocalVariableList(self):
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1(1, name="v")
+      w = variable_v1.VariableV1(
+          v,
+          trainable=False,
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
+          name="w")
+      with self.cached_session():
+        self.assertEqual(False, variable_v1.is_variable_initialized(v).eval())
+        self.assertEqual(False, variable_v1.is_variable_initialized(w).eval())
+      sm2 = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables())
+      with self.assertRaisesRegex(RuntimeError,
+                                  "Init operations did not make model ready"):
+        sm2.prepare_session("", init_op=[v.initializer])
+
+  def testPrepareSessionWithReadyNotReadyForLocal(self):
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1(1, name="v")
+      w = variable_v1.VariableV1(
+          v,
+          trainable=False,
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
+          name="w")
+      with self.cached_session():
+        self.assertEqual(False, variable_v1.is_variable_initialized(v).eval())
+        self.assertEqual(False, variable_v1.is_variable_initialized(w).eval())
+      sm2 = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables(),
+          ready_for_local_init_op=variables.report_uninitialized_variables(
+              variables.global_variables()),
+          local_init_op=w.initializer)
+      with self.assertRaisesRegex(
+          RuntimeError,
+          "Init operations did not make model ready for local_init"):
+        sm2.prepare_session("", init_op=None)
+
+  @test_util.run_v1_only("Requires TF V1 variable behavior.")
+  def testPrepareSessionWithInsufficientReadyForLocalInitCheck(self):
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1(1, name="v")
+      w = variable_v1.VariableV1(
+          v,
+          trainable=False,
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
+          name="w")
+      with self.cached_session():
+        self.assertEqual(False, variable_v1.is_variable_initialized(v).eval())
+        self.assertEqual(False, variable_v1.is_variable_initialized(w).eval())
+      sm2 = session_manager.SessionManager(
+          ready_op=variables.report_uninitialized_variables(),
+          ready_for_local_init_op=None,
+          local_init_op=w.initializer)
+    with self.assertRaisesRegex(RuntimeError,
+                                "Init operations did not make model ready.*"):
+      sm2.prepare_session("", init_op=None)
+
+
+class ObsoleteSessionManagerTest(test.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    super(ObsoleteSessionManagerTest, cls).setUpClass()
+    resource_variables_toggle.disable_resource_variables()
+
+  def testPrepareSessionSucceeds(self):
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1([1.0, 2.0, 3.0], name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.assert_variables_initialized())
+      sess = sm.prepare_session(
+          "", init_op=variables.global_variables_initializer())
+      self.assertAllClose([1.0, 2.0, 3.0], sess.run(v))
+
+  def testPrepareSessionSucceedsWithInitFeedDict(self):
+    with ops.Graph().as_default():
+      p = array_ops.placeholder(dtypes.float32, shape=(3,))
+      v = variable_v1.VariableV1(p, name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.assert_variables_initialized())
+      sess = sm.prepare_session(
+          "",
+          init_op=variables.global_variables_initializer(),
+          init_feed_dict={p: [1.0, 2.0, 3.0]})
+      self.assertAllClose([1.0, 2.0, 3.0], sess.run(v))
+
+  def testPrepareSessionSucceedsWithInitFn(self):
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1([125], name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.assert_variables_initialized())
+      sess = sm.prepare_session(
+          "", init_fn=lambda sess: sess.run(v.initializer))
+      self.assertAllClose([125], sess.run(v))
+
+  def testPrepareSessionFails(self):
+    checkpoint_dir = os.path.join(self.get_temp_dir(), "prepare_session")
+    checkpoint_dir2 = os.path.join(self.get_temp_dir(), "prepare_session2")
+    try:
+      gfile.DeleteRecursively(checkpoint_dir)
+      gfile.DeleteRecursively(checkpoint_dir2)
+    except errors.OpError:
+      pass  # Ignore
+    gfile.MakeDirs(checkpoint_dir)
+
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1([1.0, 2.0, 3.0], name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.assert_variables_initialized())
+      saver = saver_lib.Saver({"v": v})
+      sess = sm.prepare_session(
+          "",
+          init_op=variables.global_variables_initializer(),
+          saver=saver,
+          checkpoint_dir=checkpoint_dir)
+      self.assertAllClose([1.0, 2.0, 3.0], sess.run(v))
+      checkpoint_filename = os.path.join(checkpoint_dir,
+                                         "prepare_session_checkpoint")
+      saver.save(sess, checkpoint_filename)
+    # Create a new Graph and SessionManager and recover.
+    with ops.Graph().as_default():
+      # Renames the checkpoint directory.
+      os.rename(checkpoint_dir, checkpoint_dir2)
+      gfile.MakeDirs(checkpoint_dir)
+      v = variable_v1.VariableV1([6.0, 7.0, 8.0], name="v")
+      with self.cached_session():
+        self.assertEqual(False, variable_v1.is_variable_initialized(v).eval())
+      session_manager.SessionManager(
+          ready_op=variables.assert_variables_initialized())
+      saver = saver_lib.Saver({"v": v})
+      # This should fail as there's no checkpoint within 2 seconds.
+      with self.assertRaisesRegex(
+          RuntimeError, "no init_op or init_fn or local_init_op was given"):
+        sess = sm.prepare_session(
+            "",
+            init_op=None,
+            saver=saver,
+            checkpoint_dir=checkpoint_dir,
+            wait_for_checkpoint=True,
+            max_wait_secs=2)
+      # Rename the checkpoint directory back.
+      gfile.DeleteRecursively(checkpoint_dir)
+      os.rename(checkpoint_dir2, checkpoint_dir)
+      # This should succeed as there's checkpoint.
+      sess = sm.prepare_session(
+          "",
+          init_op=None,
+          saver=saver,
+          checkpoint_dir=checkpoint_dir,
+          wait_for_checkpoint=True,
+          max_wait_secs=2)
+      self.assertEqual(
+          True,
+          variable_v1.is_variable_initialized(
+              sess.graph.get_tensor_by_name("v:0")).eval(session=sess))
+
+  def testRecoverSession(self):
+    # Create a checkpoint.
+    checkpoint_dir = os.path.join(self.get_temp_dir(), "recover_session")
+    try:
+      gfile.DeleteRecursively(checkpoint_dir)
+    except errors.OpError:
+      pass  # Ignore
+    gfile.MakeDirs(checkpoint_dir)
+
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1(1, name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.assert_variables_initialized())
+      saver = saver_lib.Saver({"v": v})
+      sess, initialized = sm.recover_session(
+          "", saver=saver, checkpoint_dir=checkpoint_dir)
+      self.assertFalse(initialized)
+      sess.run(v.initializer)
+      self.assertEqual(1, sess.run(v))
+      saver.save(sess, os.path.join(checkpoint_dir,
+                                    "recover_session_checkpoint"))
+    # Create a new Graph and SessionManager and recover.
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1(2, name="v")
+      with self.cached_session():
+        self.assertEqual(False, variable_v1.is_variable_initialized(v).eval())
+      sm2 = session_manager.SessionManager(
+          ready_op=variables.assert_variables_initialized())
+      saver = saver_lib.Saver({"v": v})
+      sess, initialized = sm2.recover_session(
+          "", saver=saver, checkpoint_dir=checkpoint_dir)
+      self.assertTrue(initialized)
+      self.assertEqual(
+          True,
+          variable_v1.is_variable_initialized(
+              sess.graph.get_tensor_by_name("v:0")).eval(session=sess))
+      self.assertEqual(1, sess.run(v))
+
+  def testWaitForSessionReturnsNoneAfterTimeout(self):
+    with ops.Graph().as_default():
+      variable_v1.VariableV1(1, name="v")
+      sm = session_manager.SessionManager(
+          ready_op=variables.assert_variables_initialized(),
+          recovery_wait_secs=1)
+
+      # Set max_wait_secs to allow us to try a few times.
+      with self.assertRaises(errors.DeadlineExceededError):
+        sm.wait_for_session(master="", max_wait_secs=3)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/session_run_hook.py

@@ -0,0 +1,283 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""A SessionRunHook extends `session.run()` calls for the `MonitoredSession`.
+
+SessionRunHooks are useful to track training, report progress, request early
+stopping and more. SessionRunHooks use the observer pattern and notify at the
+following points:
+ - when a session starts being used
+ - before a call to the `session.run()`
+ - after a call to the `session.run()`
+ - when the session closed
+
+A SessionRunHook encapsulates a piece of reusable/composable computation that
+can piggyback a call to `MonitoredSession.run()`. A hook can add any
+ops-or-tensor/feeds to the run call, and when the run call finishes with success
+gets the outputs it requested. Hooks are allowed to add ops to the graph in
+`hook.begin()`. The graph is finalized after the `begin()` method is called.
+
+There are a few pre-defined hooks:
+ - StopAtStepHook: Request stop based on global_step
+ - CheckpointSaverHook: saves checkpoint
+ - LoggingTensorHook: outputs one or more tensor values to log
+ - NanTensorHook: Request stop if given `Tensor` contains Nans.
+ - SummarySaverHook: saves summaries to a summary writer
+
+For more specific needs, you can create custom hooks:
+  class ExampleHook(SessionRunHook):
+    def begin(self):
+      # You can add ops to the graph here.
+      print('Starting the session.')
+      self.your_tensor = ...
+
+    def after_create_session(self, session, coord):
+      # When this is called, the graph is finalized and
+      # ops can no longer be added to the graph.
+      print('Session created.')
+
+    def before_run(self, run_context):
+      print('Before calling session.run().')
+      return SessionRunArgs(self.your_tensor)
+
+    def after_run(self, run_context, run_values):
+      print('Done running one step. The value of my tensor: %s',
+            run_values.results)
+      if you-need-to-stop-loop:
+        run_context.request_stop()
+
+    def end(self, session):
+      print('Done with the session.')
+
+To understand how hooks interact with calls to `MonitoredSession.run()`,
+look at following code:
+  with MonitoredTrainingSession(hooks=your_hooks, ...) as sess:
+    while not sess.should_stop():
+      sess.run(your_fetches)
+
+Above user code leads to following execution:
+  call hooks.begin()
+  sess = tf.compat.v1.Session()
+  call hooks.after_create_session()
+  while not stop is requested:
+    call hooks.before_run()
+    try:
+      results = sess.run(merged_fetches, feed_dict=merged_feeds)
+    except (errors.OutOfRangeError, StopIteration):
+      break
+    call hooks.after_run()
+  call hooks.end()
+  sess.close()
+
+Note that if sess.run() raises OutOfRangeError or StopIteration then
+hooks.after_run() will not be called but hooks.end() will still be called.
+If sess.run() raises any other exception then neither hooks.after_run() nor
+hooks.end() will be called.
+"""
+
+import collections
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(v1=["train.SessionRunHook"])
+class SessionRunHook:
+  """Hook to extend calls to MonitoredSession.run()."""
+
+  def begin(self):
+    """Called once before using the session.
+
+    When called, the default graph is the one that will be launched in the
+    session.  The hook can modify the graph by adding new operations to it.
+    After the `begin()` call the graph will be finalized and the other callbacks
+    can not modify the graph anymore. Second call of `begin()` on the same
+    graph, should not change the graph.
+    """
+    pass
+
+  def after_create_session(self, session, coord):  # pylint: disable=unused-argument
+    """Called when new TensorFlow session is created.
+
+    This is called to signal the hooks that a new session has been created. This
+    has two essential differences with the situation in which `begin` is called:
+
+    * When this is called, the graph is finalized and ops can no longer be added
+        to the graph.
+    * This method will also be called as a result of recovering a wrapped
+        session, not only at the beginning of the overall session.
+
+    Args:
+      session: A TensorFlow Session that has been created.
+      coord: A Coordinator object which keeps track of all threads.
+    """
+    pass
+
+  def before_run(self, run_context):  # pylint: disable=unused-argument
+    """Called before each call to run().
+
+    You can return from this call a `SessionRunArgs` object indicating ops or
+    tensors to add to the upcoming `run()` call.  These ops/tensors will be run
+    together with the ops/tensors originally passed to the original run() call.
+    The run args you return can also contain feeds to be added to the run()
+    call.
+
+    The `run_context` argument is a `SessionRunContext` that provides
+    information about the upcoming `run()` call: the originally requested
+    op/tensors, the TensorFlow Session.
+
+    At this point graph is finalized and you can not add ops.
+
+    Args:
+      run_context: A `SessionRunContext` object.
+
+    Returns:
+      None or a `SessionRunArgs` object.
+    """
+    return None
+
+  def after_run(self,
+                run_context,  # pylint: disable=unused-argument
+                run_values):  # pylint: disable=unused-argument
+    """Called after each call to run().
+
+    The `run_values` argument contains results of requested ops/tensors by
+    `before_run()`.
+
+    The `run_context` argument is the same one send to `before_run` call.
+    `run_context.request_stop()` can be called to stop the iteration.
+
+    If `session.run()` raises any exceptions then `after_run()` is not called.
+
+    Args:
+      run_context: A `SessionRunContext` object.
+      run_values: A SessionRunValues object.
+    """
+    pass
+
+  def end(self, session):  # pylint: disable=unused-argument
+    """Called at the end of session.
+
+    The `session` argument can be used in case the hook wants to run final ops,
+    such as saving a last checkpoint.
+
+    If `session.run()` raises exception other than OutOfRangeError or
+    StopIteration then `end()` is not called.
+    Note the difference between `end()` and `after_run()` behavior when
+    `session.run()` raises OutOfRangeError or StopIteration. In that case
+    `end()` is called but `after_run()` is not called.
+
+    Args:
+      session: A TensorFlow Session that will be soon closed.
+    """
+    pass
+
+
+@tf_export(v1=["train.SessionRunArgs"])
+class SessionRunArgs(
+    collections.namedtuple("SessionRunArgs",
+                           ["fetches", "feed_dict", "options"])):
+  """Represents arguments to be added to a `Session.run()` call.
+
+  Args:
+    fetches: Exactly like the 'fetches' argument to Session.Run().
+      Can be a single tensor or op, a list of 'fetches' or a dictionary
+      of fetches.  For example:
+        fetches = global_step_tensor
+        fetches = [train_op, summary_op, global_step_tensor]
+        fetches = {'step': global_step_tensor, 'summ': summary_op}
+      Note that this can recurse as expected:
+        fetches = {'step': global_step_tensor,
+                   'ops': [train_op, check_nan_op]}
+    feed_dict: Exactly like the `feed_dict` argument to `Session.Run()`
+    options: Exactly like the `options` argument to `Session.run()`, i.e., a
+      config_pb2.RunOptions proto.
+  """
+
+  def __new__(cls, fetches, feed_dict=None, options=None):
+    return super(SessionRunArgs, cls).__new__(cls, fetches, feed_dict, options)
+
+
+@tf_export(v1=["train.SessionRunContext"])
+class SessionRunContext:
+  """Provides information about the `session.run()` call being made.
+
+  Provides information about original request to `Session.Run()` function.
+  SessionRunHook objects can stop the loop by calling `request_stop()` of
+  `run_context`. In the future we may use this object to add more information
+  about run without changing the Hook API.
+  """
+
+  def __init__(self, original_args, session):
+    """Initializes SessionRunContext."""
+    self._original_args = original_args
+    self._session = session
+    self._stop_requested = False
+
+  @property
+  def original_args(self):
+    """A `SessionRunArgs` object holding the original arguments of `run()`.
+
+    If user called `MonitoredSession.run(fetches=a, feed_dict=b)`, then this
+    field is equal to SessionRunArgs(a, b).
+
+    Returns:
+     A `SessionRunArgs` object
+    """
+    return self._original_args
+
+  @property
+  def session(self):
+    """A TensorFlow session object which will execute the `run`."""
+    return self._session
+
+  @property
+  def stop_requested(self):
+    """Returns whether a stop is requested or not.
+
+    If true, `MonitoredSession` stops iterations.
+    Returns:
+      A `bool`
+    """
+    return self._stop_requested
+
+  def request_stop(self):
+    """Sets stop requested field.
+
+    Hooks can use this function to request stop of iterations.
+    `MonitoredSession` checks whether this is called or not.
+    """
+    self._stop_requested = True
+
+
+@tf_export(v1=["train.SessionRunValues"])
+class SessionRunValues(
+    collections.namedtuple("SessionRunValues",
+                           ["results", "options", "run_metadata"])):
+  """Contains the results of `Session.run()`.
+
+  In the future we may use this object to add more information about result of
+  run without changing the Hook API.
+
+  Args:
+    results: The return values from `Session.run()` corresponding to the fetches
+      attribute returned in the RunArgs. Note that this has the same shape as
+      the RunArgs fetches.  For example:
+        fetches = global_step_tensor
+        => results = nparray(int)
+        fetches = [train_op, summary_op, global_step_tensor]
+        => results = [None, nparray(string), nparray(int)]
+        fetches = {'step': global_step_tensor, 'summ': summary_op}
+        => results = {'step': nparray(int), 'summ': nparray(string)}
+    options: `RunOptions` from the `Session.run()` call.
+    run_metadata: `RunMetadata` from the `Session.run()` call.
+  """

tensorflow/python/training/slot_creator.py

@@ -0,0 +1,281 @@
+# 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.
+# ==============================================================================
+
+"""Standard functions for creating slots.
+
+A slot is a `Variable` created with the same first m-dimension as a primary
+variable or `Tensor`. A slot is always scoped in the namespace of the primary
+object and typically has the same device and type.
+
+Slots are typically used as accumulators to track values associated with
+the primary object:
+
+```python
+# Optimizers can create a slot for each variable to track accumulators
+accumulators = {var : create_zeros_slot(var, "momentum") for var in vs}
+for var in vs:
+  apply_momentum(var, accumulators[var], lr, grad, momentum_tensor)
+
+# Slots can also be used for moving averages
+mavg = create_slot(var, var.initialized_value(), "exponential_moving_avg")
+update_mavg = mavg.assign_sub((mavg - var) * (1 - decay))
+```
+"""
+# pylint: disable=g-bad-name
+
+from tensorflow.python.compiler.xla.experimental import xla_sharding
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import ref_variable
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.ops import variables
+
+
+def _create_slot_var(primary,
+                     val,
+                     scope,
+                     validate_shape,
+                     shape,
+                     dtype,
+                     *,
+                     copy_xla_sharding=False):
+  """Helper function for creating a slot variable."""
+
+  # TODO(lukaszkaiser): Consider allowing partitioners to be set in the current
+  # scope.
+  current_partitioner = variable_scope.get_variable_scope().partitioner
+  variable_scope.get_variable_scope().set_partitioner(None)
+  # When init from val instead of callable initializer, the shape is expected to
+  # be None, not <unknown> or any fully defined shape.
+  shape = shape if callable(val) else None
+  if resource_variable_ops.is_resource_variable(primary):
+    use_resource = True
+  elif isinstance(primary, ref_variable.RefVariable):
+    use_resource = False
+  else:
+    use_resource = None
+  slot = variable_scope.get_variable(
+      scope,
+      initializer=val,
+      trainable=False,
+      use_resource=use_resource,
+      shape=shape,
+      dtype=dtype,
+      validate_shape=validate_shape)
+  variable_scope.get_variable_scope().set_partitioner(current_partitioner)
+
+  # pylint: disable=protected-access
+  if isinstance(primary, variables.Variable) and primary._save_slice_info:
+    # Primary is a partitioned variable, so we need to also indicate that
+    # the slot is a partitioned variable.  Slots have the same partitioning
+    # as their primaries.
+    # For examples when using AdamOptimizer in linear model, slot.name
+    # here can be "linear//weights/Adam:0", while primary.op.name is
+    # "linear//weight". We want to get 'Adam' as real_slot_name, so we
+    # remove "'linear//weight' + '/'" and ':0'.
+    real_slot_name = slot.name[len(primary.op.name + "/"):-2]
+    slice_info = primary._save_slice_info
+    # support slot's shape not same as primary's shape
+    # example: primary's shape = [10, 20, 30], slot's shape =
+    # None, [], [10], [10, 20] or [10, 20, 30] is allowed
+    # slot's shape = None or [10, 20, 30], set slot's slice_info same as primary
+    # slot's shape = [], don't set slot's slice_info
+    # slot's shape = [10] or [10, 20], set slot's slice_info according to ndims
+    n = slot.shape.ndims
+    if n is None or n > 0:
+      slot._set_save_slice_info(
+          variables.Variable.SaveSliceInfo(
+              slice_info.full_name + "/" + real_slot_name,
+              slice_info.full_shape[:n], slice_info.var_offset[:n],
+              slice_info.var_shape[:n]))
+  # pylint: enable=protected-access
+
+  # Copy XLA sharding attributes from the primary if the slot variable has the
+  # same rank as the primary.
+  def _has_same_rank(primary_shape, slot_shape):
+    return (primary_shape.rank is not None and slot_shape.rank is not None and
+            primary_shape.rank == slot_shape.rank)
+
+  if copy_xla_sharding and _has_same_rank(primary.shape, slot.shape):
+    slot = xla_sharding.copy_sharding(primary, slot, use_sharding_op=False)
+  return slot
+
+
+def create_slot(primary,
+                val,
+                name,
+                colocate_with_primary=True,
+                *,
+                copy_xla_sharding=False):
+  """Create a slot initialized to the given value.
+
+  The type of the slot is determined by the given value.
+
+  Args:
+    primary: The primary `Variable` or `Tensor`.
+    val: A `Tensor` specifying the initial value of the slot.
+    name: Name to use for the slot variable.
+    colocate_with_primary: Boolean.  If True the slot is located
+      on the same device as `primary`.
+    copy_xla_sharding: Boolean. If True also copies XLA sharding
+      from primary.
+
+  Returns:
+    A `Variable` object.
+  """
+  # Scope the slot name in the namespace of the primary variable.
+  # Set primary's name + '/' + name as default name, so the scope name of
+  # optimizer can be shared when reuse is True. Meanwhile when reuse is False
+  # and the same name has been previously used, the scope name will add '_N'
+  # as suffix for unique identifications.
+  validate_shape = val.get_shape().is_fully_defined()
+  if isinstance(primary, variables.Variable):
+    prefix = primary._shared_name  # pylint: disable=protected-access
+  else:
+    prefix = primary.op.name
+  with variable_scope.variable_scope(None, prefix + "/" + name):
+    if colocate_with_primary:
+      distribution_strategy = distribute_lib.get_strategy()
+      with distribution_strategy.extended.colocate_vars_with(primary):
+        return _create_slot_var(
+            primary,
+            val,
+            "",
+            validate_shape,
+            None,
+            None,
+            copy_xla_sharding=copy_xla_sharding)
+    else:
+      return _create_slot_var(
+          primary,
+          val,
+          "",
+          validate_shape,
+          None,
+          None,
+          copy_xla_sharding=copy_xla_sharding)
+
+
+def create_slot_with_initializer(primary,
+                                 initializer,
+                                 shape,
+                                 dtype,
+                                 name,
+                                 colocate_with_primary=True,
+                                 *,
+                                 copy_xla_sharding=False):
+  """Creates a slot initialized using an `Initializer`.
+
+  The type of the slot is determined by the given value.
+
+  Args:
+    primary: The primary `Variable` or `Tensor`.
+    initializer: An `Initializer`.  The initial value of the slot.
+    shape: Shape of the initial value of the slot.
+    dtype: Type of the value of the slot.
+    name: Name to use for the slot variable.
+    colocate_with_primary: Boolean.  If True the slot is located
+      on the same device as `primary`.
+    copy_xla_sharding: Boolean. If True also copies XLA sharding
+      from primary.
+
+  Returns:
+    A `Variable` object.
+  """
+  # Scope the slot name in the namespace of the primary variable.
+  # Set "primary.op.name + '/' + name" as default name, so the scope name of
+  # optimizer can be shared when reuse is True. Meanwhile when reuse is False
+  # and the same name has been previously used, the scope name will add '_N'
+  # as suffix for unique identifications.
+  validate_shape = shape.is_fully_defined()
+  if isinstance(primary, variables.Variable):
+    prefix = primary._shared_name  # pylint: disable=protected-access
+  else:
+    prefix = primary.op.name
+  with variable_scope.variable_scope(None, prefix + "/" + name):
+    if colocate_with_primary:
+      distribution_strategy = distribute_lib.get_strategy()
+      with distribution_strategy.extended.colocate_vars_with(primary):
+        return _create_slot_var(
+            primary,
+            initializer,
+            "",
+            validate_shape,
+            shape,
+            dtype,
+            copy_xla_sharding=copy_xla_sharding)
+    else:
+      return _create_slot_var(
+          primary,
+          initializer,
+          "",
+          validate_shape,
+          shape,
+          dtype,
+          copy_xla_sharding=copy_xla_sharding)
+
+
+def create_zeros_slot(primary,
+                      name,
+                      dtype=None,
+                      colocate_with_primary=True,
+                      *,
+                      copy_xla_sharding=False):
+  """Create a slot initialized to 0 with same shape as the primary object.
+
+  Args:
+    primary: The primary `Variable` or `Tensor`.
+    name: Name to use for the slot variable.
+    dtype: Type of the slot variable.  Defaults to the type of `primary`.
+    colocate_with_primary: Boolean.  If True the slot is located
+      on the same device as `primary`.
+    copy_xla_sharding: Boolean. If True also copies XLA sharding
+      from primary.
+
+  Returns:
+    A `Variable` object.
+  """
+  if dtype is None:
+    dtype = primary.dtype
+  slot_shape = primary.get_shape()
+  if slot_shape.is_fully_defined():
+    initializer = init_ops.zeros_initializer()
+    return create_slot_with_initializer(
+        primary,
+        initializer,
+        slot_shape,
+        dtype,
+        name,
+        colocate_with_primary=colocate_with_primary,
+        copy_xla_sharding=copy_xla_sharding)
+  else:
+    if isinstance(primary, variables.Variable):
+      slot_shape = array_ops.shape(
+          cond.cond(
+              variable_v1.is_variable_initialized(primary), primary.read_value,
+              lambda: primary.initial_value))
+    else:
+      slot_shape = array_ops.shape(primary)
+    val = array_ops.zeros(slot_shape, dtype=dtype)
+    return create_slot(
+        primary,
+        val,
+        name,
+        colocate_with_primary=colocate_with_primary,
+        copy_xla_sharding=copy_xla_sharding)

tensorflow/python/training/slot_creator_test.py

@@ -0,0 +1,283 @@
+# 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 slot_creator."""
+
+import numpy as np
+
+from xla import xla_data_pb2
+from tensorflow.python.compiler.xla.experimental import xla_sharding
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import partitioned_variables
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import slot_creator
+
+
+def initialized_value(var):
+  return cond.cond(
+      variable_v1.is_variable_initialized(var), var.read_value,
+      lambda: var.initial_value)
+
+
+class SlotCreatorTest(test.TestCase):
+
+  def testCreateSlotFromVariable(self):
+    # slot_creator is used only in optimizer V1.
+    with ops.Graph().as_default(), self.cached_session():
+      v = variables.Variable([1.0, 2.5], name="var")
+      slot = slot_creator.create_slot(v, initialized_value(v), name="slot")
+
+      self.evaluate(variables.global_variables_initializer())
+
+      self.assertEqual("var/slot", slot.op.name)
+      self.assertEqual([2], slot.get_shape().as_list())
+      self.assertEqual(dtypes.float32, slot.dtype.base_dtype)
+      self.assertAllEqual([1.0, 2.5], self.evaluate(slot))
+
+  def testCreateSlotFromTensor(self):
+    # slot_creator is used only in optimizer V1.
+    with ops.Graph().as_default(), self.cached_session():
+      v = constant_op.constant([1.0, 2.5], name="const")
+      slot = slot_creator.create_slot(v, v * 2, name="slot")
+
+      self.evaluate(variables.global_variables_initializer())
+
+      self.assertEqual("const/slot", slot.op.name)
+      self.assertEqual([2], slot.get_shape().as_list())
+      self.assertEqual(dtypes.float32, slot.dtype.base_dtype)
+      self.assertAllEqual([2.0, 5.0], self.evaluate(slot))
+
+  def testCreateZerosSlotFromVariable(self):
+    # slot_creator is used only in optimizer V1.
+    with ops.Graph().as_default(), self.cached_session():
+      v = variables.Variable([1.0, 2.5], name="var")
+      with ops.control_dependencies(None):
+        slot = slot_creator.create_zeros_slot(
+            v, name="slot", dtype=dtypes.float64)
+
+      self.evaluate(variables.global_variables_initializer())
+
+      self.assertEqual("var/slot", slot.op.name)
+      self.assertEqual([2], slot.get_shape().as_list())
+      self.assertEqual(dtypes.float64, slot.dtype.base_dtype)
+      self.assertAllEqual([0.0, 0.0], self.evaluate(slot))
+
+  def testCreateZerosSlotFromDynamicShapedVariable(self):
+    # slot_creator is used only in optimizer V1.
+    with ops.Graph().as_default(), self.cached_session():
+      dyn_shape = constant_op.constant([2], dtype=dtypes.int32)
+      dyn_shape = array_ops.placeholder_with_default(dyn_shape,
+                                                     shape=[None])
+      v = variable_scope.get_variable(
+          "var",
+          initializer=random_ops.random_uniform(dyn_shape,
+                                                dtype=dtypes.float64),
+          validate_shape=False)
+      with ops.control_dependencies(None):
+        slot = slot_creator.create_zeros_slot(
+            v, name="slot", dtype=dtypes.float64)
+
+      self.evaluate(variables.global_variables_initializer())
+
+      self.assertEqual("var/slot", slot.op.name)
+      self.assertEqual([2], array_ops.shape(slot).eval())
+      self.assertEqual(dtypes.float64, slot.dtype.base_dtype)
+      self.assertAllEqual([0.0, 0.0], self.evaluate(slot))
+
+  def testCreateZerosSlotFromTensor(self):
+    # slot_creator is used only in optimizer V1.
+    with ops.Graph().as_default(), self.cached_session():
+      v = constant_op.constant([1.0, 2.5], name="const")
+      with ops.control_dependencies(None):
+        slot = slot_creator.create_zeros_slot(v, name="slot")
+
+      self.evaluate(variables.global_variables_initializer())
+
+      self.assertEqual("const/slot", slot.op.name)
+      self.assertEqual([2], slot.get_shape().as_list())
+      self.assertEqual(dtypes.float32, slot.dtype.base_dtype)
+      self.assertAllEqual([0.0, 0.0], self.evaluate(slot))
+
+  def testCreateZerosSlotFromDynamicShapedTensor(self):
+    # slot_creator is used only in optimizer V1.
+    with ops.Graph().as_default(), self.cached_session():
+      v = random_ops.random_uniform([2], dtype=dtypes.float64)
+      v = array_ops.placeholder_with_default(v, shape=[None], name="const")
+      with ops.control_dependencies(None):
+        slot = slot_creator.create_zeros_slot(
+            v, name="slot", dtype=dtypes.float64)
+
+      self.evaluate(variables.global_variables_initializer())
+
+      self.assertEqual("const/slot", slot.op.name)
+      self.assertEqual([2], array_ops.shape(slot).eval())
+      self.assertEqual(dtypes.float64, slot.dtype.base_dtype)
+      self.assertAllEqual([0.0, 0.0], self.evaluate(slot))
+
+  def testCreateSlotFromVariableRespectsScope(self):
+    # See discussion on #2740.
+    # slot_creator is used only in optimizer V1.
+    with ops.Graph().as_default(), self.cached_session():
+      with variable_scope.variable_scope("scope"):
+        v = variables.Variable([1.0, 2.5], name="var")
+        slot = slot_creator.create_slot(v, initialized_value(v), name="slot")
+        self.assertEqual("scope/scope/var/slot", slot.op.name)
+
+  def testCreateSlotFromFirstMDimensionVariable(self):
+    # slot_creator is used only in optimizer V1.
+    with ops.Graph().as_default(), self.test_session():
+      s = variables.Variable([1.0, 2.5], name="var")
+      p_v = variable_scope.get_variable(
+          "var",
+          shape=[2, 2],
+          partitioner=partitioned_variables.fixed_size_partitioner(2))
+      for i, v in enumerate(p_v):
+        slot = slot_creator.create_slot(v, initialized_value(s), name="slot")
+        si = slot._save_slice_info
+
+        self.evaluate(variables.global_variables_initializer())
+
+        self.assertEqual("var/part_%d/slot" % i, slot.op.name)
+        self.assertEqual([2], slot.get_shape().as_list())
+        self.assertEqual(dtypes.float32, slot.dtype.base_dtype)
+        self.assertAllEqual([1.0, 2.5], slot)
+        self.assertAllEqual([2], si.full_shape)
+        self.assertAllEqual([i], si.var_offset)
+        self.assertAllEqual([1], si.var_shape)
+
+  def testCreateSlotFromScalarVariable(self):
+    # slot_creator is used only in optimizer V1.
+    with ops.Graph().as_default(), self.test_session():
+      s = variables.Variable(1.0, name="var")
+      p_v = variable_scope.get_variable(
+          "var",
+          shape=[2, 2],
+          partitioner=partitioned_variables.fixed_size_partitioner(2))
+      for i, v in enumerate(p_v):
+        slot = slot_creator.create_slot(v, initialized_value(s), name="slot")
+
+        self.evaluate(variables.global_variables_initializer())
+
+        self.assertEqual("var/part_%d/slot" % i, slot.op.name)
+        self.assertEqual([], slot.get_shape().as_list())
+        self.assertEqual(dtypes.float32, slot.dtype.base_dtype)
+        self.assertAllEqual(1.0, slot)
+
+  def testCreateSlotFromVariableCopyXlaSharding(self):
+    # slot_creator is used only in optimizer V1.
+    with ops.Graph().as_default(), self.cached_session():
+      v = variables.Variable([1.0, 2.5], name="var")
+      v = xla_sharding.mesh_split(
+          v, np.array([0, 1]), [0], use_sharding_op=False)
+      slot = slot_creator.create_slot(
+          v, initialized_value(v), name="slot", copy_xla_sharding=True)
+      self.assertEqual(
+          xla_sharding.get_tensor_sharding(v),
+          xla_sharding.get_tensor_sharding(slot))
+
+  def testCreateZerosSlotFromVariableCopyXlaSharding(self):
+    # slot_creator is used only in optimizer V1.
+    with ops.Graph().as_default(), self.cached_session():
+      v = variables.Variable([1.0, 2.5], name="var")
+      v = xla_sharding.mesh_split(
+          v, np.array([0, 1]), [0], use_sharding_op=False)
+      with ops.control_dependencies(None):
+        slot = slot_creator.create_zeros_slot(
+            v, name="slot", dtype=dtypes.float64, copy_xla_sharding=True)
+      self.assertEqual(
+          xla_sharding.get_tensor_sharding(v),
+          xla_sharding.get_tensor_sharding(slot))
+
+  def testCreateSlotWithoutXlaSharding(self):
+    # slot_creator is used only in optimizer V1.
+    # The SPMD sharding annotations should not be copied since the primary
+    # variable and slot variable have different ranks.
+    with ops.Graph().as_default(), self.cached_session():
+      v = variables.Variable([1.0, 2.5], name="var")
+      v = xla_sharding.mesh_split(
+          v, np.array([0, 1]), [0], use_sharding_op=False)
+      with ops.control_dependencies(None):
+        slot = slot_creator.create_slot(
+            v,
+            constant_op.constant(10, name="const"),
+            name="slot",
+            copy_xla_sharding=True)
+      self.assertIsNone(xla_sharding.get_tensor_sharding(slot))
+      self.assertNotEqual(
+          xla_sharding.get_tensor_sharding(v),
+          xla_sharding.get_tensor_sharding(slot))
+
+  def testCreateSlotWithCustomReplicatedXlaSharding(self):
+    # slot_creator is used only in optimizer V1.
+    # We insert our own custom replicated XLA sharding that overrides the SPMD
+    # sharding copied over by the slot_creator.
+    with ops.Graph().as_default(), self.cached_session():
+      v = variables.Variable([1.0, 2.5], name="var")
+      v = xla_sharding.mesh_split(
+          v, np.array([0, 1]), [0], use_sharding_op=False)
+      with ops.control_dependencies(None):
+        slot = slot_creator.create_zeros_slot(
+            v, name="slot", dtype=dtypes.float64, copy_xla_sharding=True)
+        slot = xla_sharding.replicate(slot, use_sharding_op=False)
+
+      self.assertNotEqual(
+          xla_sharding.get_tensor_sharding(v),
+          xla_sharding.get_tensor_sharding(slot))
+
+      slot_sharding = xla_sharding.get_tensor_sharding(slot)
+      slot_proto = xla_data_pb2.OpSharding()
+      slot_proto.ParseFromString(slot_sharding)
+      self.assertEqual(
+          slot_proto,
+          xla_data_pb2.OpSharding(type=xla_data_pb2.OpSharding.REPLICATED))
+
+  def testCreateSlotWithCustomSplitXlaSharding(self):
+    # slot_creator is used only in optimizer V1.
+    # We insert our own custom split XLA sharding that overrides the SPMD
+    # sharding copied over by the slot_creator.
+    with ops.Graph().as_default(), self.cached_session():
+      v = variables.Variable([1.0, 2.5, 10.0, 15.1], name="var")
+      v = xla_sharding.mesh_split(
+          v, np.array([0, 1]), [0], use_sharding_op=False)
+      with ops.control_dependencies(None):
+        slot = slot_creator.create_zeros_slot(
+            v, name="slot", dtype=dtypes.float64, copy_xla_sharding=True)
+        slot = xla_sharding.split(
+            slot, split_dimension=0, num_devices=4, use_sharding_op=False)
+
+      self.assertNotEqual(
+          xla_sharding.get_tensor_sharding(v),
+          xla_sharding.get_tensor_sharding(slot))
+
+      slot_sharding = xla_sharding.get_tensor_sharding(slot)
+      slot_proto = xla_data_pb2.OpSharding()
+      slot_proto.ParseFromString(slot_sharding)
+      self.assertEqual(
+          slot_proto,
+          xla_data_pb2.OpSharding(
+              type=xla_data_pb2.OpSharding.OTHER,
+              tile_assignment_dimensions=[4],
+              tile_assignment_devices=range(4)))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/summary_io.py

@@ -0,0 +1,77 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Reads Summaries from and writes Summaries to event files."""
+
+# pylint: disable=unused-import
+from tensorflow.python.summary.summary_iterator import summary_iterator
+from tensorflow.python.summary.writer.writer import FileWriter as _FileWriter
+from tensorflow.python.summary.writer.writer_cache import FileWriterCache as SummaryWriterCache
+# pylint: enable=unused-import
+from tensorflow.python.util.deprecation import deprecated
+
+
+class SummaryWriter(_FileWriter):
+
+  @deprecated("2016-11-30",
+              "Please switch to tf.summary.FileWriter. The interface and "
+              "behavior is the same; this is just a rename.")
+  def __init__(self,
+               logdir,
+               graph=None,
+               max_queue=10,
+               flush_secs=120,
+               graph_def=None):
+    """Creates a `SummaryWriter` and an event file.
+
+    This class is deprecated, and should be replaced with tf.summary.FileWriter.
+
+    On construction the summary writer creates a new event file in `logdir`.
+    This event file will contain `Event` protocol buffers constructed when you
+    call one of the following functions: `add_summary()`, `add_session_log()`,
+    `add_event()`, or `add_graph()`.
+
+    If you pass a `Graph` to the constructor it is added to
+    the event file. (This is equivalent to calling `add_graph()` later).
+
+    TensorBoard will pick the graph from the file and display it graphically so
+    you can interactively explore the graph you built. You will usually pass
+    the graph from the session in which you launched it:
+
+    ```python
+    ...create a graph...
+    # Launch the graph in a session.
+    sess = tf.compat.v1.Session()
+    # Create a summary writer, add the 'graph' to the event file.
+    writer = tf.compat.v1.summary.FileWriter(<some-directory>, sess.graph)
+    ```
+
+    The other arguments to the constructor control the asynchronous writes to
+    the event file:
+
+    *  `flush_secs`: How often, in seconds, to flush the added summaries
+       and events to disk.
+    *  `max_queue`: Maximum number of summaries or events pending to be
+       written to disk before one of the 'add' calls block.
+
+    Args:
+      logdir: A string. Directory where event file will be written.
+      graph: A `Graph` object, such as `sess.graph`.
+      max_queue: Integer. Size of the queue for pending events and summaries.
+      flush_secs: Number. How often, in seconds, to flush the
+        pending events and summaries to disk.
+      graph_def: DEPRECATED: Use the `graph` argument instead.
+    """
+    super(SummaryWriter, self).__init__(logdir, graph, max_queue, flush_secs,
+                                        graph_def)

tensorflow/python/training/supervisor_test.py

@@ -0,0 +1,916 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Tests for supervisor.py."""
+
+import glob
+import os
+import shutil
+import time
+import uuid
+
+
+from tensorflow.core.framework import graph_pb2
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.core.protobuf import meta_graph_pb2
+from tensorflow.core.util import event_pb2
+from tensorflow.python.checkpoint import checkpoint_management
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.framework import meta_graph
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import io_ops
+from tensorflow.python.ops import parsing_ops
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import gfile
+from tensorflow.python.platform import test
+from tensorflow.python.summary import summary
+from tensorflow.python.summary import summary_iterator
+from tensorflow.python.summary.writer import writer
+from tensorflow.python.training import input as input_lib
+from tensorflow.python.training import saver as saver_lib
+from tensorflow.python.training import server_lib
+from tensorflow.python.training import session_manager as session_manager_lib
+from tensorflow.python.training import supervisor
+
+
+def _summary_iterator(test_dir):
+  """Reads events from test_dir/events.
+
+  Args:
+    test_dir: Name of the test directory.
+
+  Returns:
+    A summary_iterator
+  """
+  event_paths = sorted(glob.glob(os.path.join(test_dir, "event*")))
+  return summary_iterator.summary_iterator(event_paths[-1])
+
+
+class SupervisorTest(test.TestCase):
+
+  def _test_dir(self, test_name):
+    test_dir = os.path.join(self.get_temp_dir(), test_name)
+    if os.path.exists(test_dir):
+      shutil.rmtree(test_dir)
+    return test_dir
+
+  def _wait_for_glob(self, pattern, timeout_secs, for_checkpoint=True):
+    """Wait for a checkpoint file to appear.
+
+    Args:
+      pattern: A string.
+      timeout_secs: How long to wait for in seconds.
+      for_checkpoint: whether we're globbing for checkpoints.
+    """
+    end_time = time.time() + timeout_secs
+    while time.time() < end_time:
+      if for_checkpoint:
+        if checkpoint_management.checkpoint_exists(pattern):
+          return
+      else:
+        if len(gfile.Glob(pattern)) >= 1:
+          return
+      time.sleep(0.05)
+    self.assertFalse(True, "Glob never matched any file: %s" % pattern)
+
+  # This test does not test much.
+  def testBasics(self):
+    logdir = self._test_dir("basics")
+    with ops.Graph().as_default():
+      my_op = constant_op.constant(1.0)
+      sv = supervisor.Supervisor(logdir=logdir)
+      sess = sv.prepare_or_wait_for_session("")
+      for _ in range(10):
+        self.evaluate(my_op)
+      sess.close()
+      sv.stop()
+
+  def testManagedSession(self):
+    logdir = self._test_dir("managed_session")
+    with ops.Graph().as_default():
+      my_op = constant_op.constant(1.0)
+      sv = supervisor.Supervisor(logdir=logdir)
+      with sv.managed_session(""):
+        for _ in range(10):
+          self.evaluate(my_op)
+      # Supervisor has been stopped.
+      self.assertTrue(sv.should_stop())
+
+  def testManagedSessionUserError(self):
+    logdir = self._test_dir("managed_user_error")
+    with ops.Graph().as_default():
+      my_op = constant_op.constant(1.0)
+      sv = supervisor.Supervisor(logdir=logdir)
+      last_step = None
+      with self.assertRaisesRegex(RuntimeError, "failing here"):
+        with sv.managed_session("") as sess:
+          for step in range(10):
+            last_step = step
+            if step == 1:
+              raise RuntimeError("failing here")
+            else:
+              self.evaluate(my_op)
+      # Supervisor has been stopped.
+      self.assertTrue(sv.should_stop())
+      self.assertEqual(1, last_step)
+
+  def testManagedSessionIgnoreOutOfRangeError(self):
+    logdir = self._test_dir("managed_out_of_range")
+    with ops.Graph().as_default():
+      my_op = constant_op.constant(1.0)
+      sv = supervisor.Supervisor(logdir=logdir)
+      last_step = None
+      with sv.managed_session("") as sess:
+        for step in range(10):
+          last_step = step
+          if step == 3:
+            raise errors_impl.OutOfRangeError(my_op.op.node_def, my_op.op,
+                                              "all done")
+          else:
+            self.evaluate(my_op)
+      # Supervisor has been stopped.  OutOfRangeError was not thrown.
+      self.assertTrue(sv.should_stop())
+      self.assertEqual(3, last_step)
+
+  def testManagedSessionDoNotKeepSummaryWriter(self):
+    logdir = self._test_dir("managed_not_keep_summary_writer")
+    with ops.Graph().as_default():
+      summary.scalar("c1", constant_op.constant(1))
+      summary.scalar("c2", constant_op.constant(2))
+      summary.scalar("c3", constant_op.constant(3))
+      summ = summary.merge_all()
+      sv = supervisor.Supervisor(logdir=logdir, summary_op=None)
+      with sv.managed_session(
+          "", close_summary_writer=True, start_standard_services=False) as sess:
+        sv.summary_computed(sess, sess.run(summ))
+      # Sleep 1.2s to make sure that the next event file has a different name
+      # than the current one.
+      time.sleep(1.2)
+      with sv.managed_session(
+          "", close_summary_writer=True, start_standard_services=False) as sess:
+        sv.summary_computed(sess, sess.run(summ))
+    event_paths = sorted(glob.glob(os.path.join(logdir, "event*")))
+    self.assertEqual(2, len(event_paths))
+    # The two event files should have the same contents.
+    for path in event_paths:
+      # The summary iterator should report the summary once as we closed the
+      # summary writer across the 2 sessions.
+      rr = summary_iterator.summary_iterator(path)
+      # The first event should list the file_version.
+      ev = next(rr)
+      self.assertEqual("brain.Event:2", ev.file_version)
+
+      # The next one has the graph and metagraph.
+      ev = next(rr)
+      self.assertTrue(ev.graph_def)
+
+      ev = next(rr)
+      self.assertTrue(ev.meta_graph_def)
+
+      # The next one should have the values from the summary.
+      # But only once.
+      ev = next(rr)
+      self.assertProtoEquals("""
+        value { tag: 'c1' simple_value: 1.0 }
+        value { tag: 'c2' simple_value: 2.0 }
+        value { tag: 'c3' simple_value: 3.0 }
+        """, ev.summary)
+
+      # The next one should be a stop message if we closed cleanly.
+      ev = next(rr)
+      self.assertEqual(event_pb2.SessionLog.STOP, ev.session_log.status)
+
+      # We should be done.
+      with self.assertRaises(StopIteration):
+        next(rr)
+
+  def testManagedSessionKeepSummaryWriter(self):
+    logdir = self._test_dir("managed_keep_summary_writer")
+    with ops.Graph().as_default():
+      summary.scalar("c1", constant_op.constant(1))
+      summary.scalar("c2", constant_op.constant(2))
+      summary.scalar("c3", constant_op.constant(3))
+      summ = summary.merge_all()
+      sv = supervisor.Supervisor(logdir=logdir)
+      with sv.managed_session(
+          "", close_summary_writer=False,
+          start_standard_services=False) as sess:
+        sv.summary_computed(sess, sess.run(summ))
+      with sv.managed_session(
+          "", close_summary_writer=False,
+          start_standard_services=False) as sess:
+        sv.summary_computed(sess, sess.run(summ))
+    # Now close the summary writer to flush the events.
+    sv.summary_writer.close()
+    # The summary iterator should report the summary twice as we reused
+    # the same summary writer across the 2 sessions.
+    rr = _summary_iterator(logdir)
+    # The first event should list the file_version.
+    ev = next(rr)
+    self.assertEqual("brain.Event:2", ev.file_version)
+
+    # The next one has the graph.
+    ev = next(rr)
+    self.assertTrue(ev.graph_def)
+
+    ev = next(rr)
+    self.assertTrue(ev.meta_graph_def)
+
+    # The next one should have the values from the summary.
+    ev = next(rr)
+    self.assertProtoEquals("""
+      value { tag: 'c1' simple_value: 1.0 }
+      value { tag: 'c2' simple_value: 2.0 }
+      value { tag: 'c3' simple_value: 3.0 }
+      """, ev.summary)
+
+    # The next one should also have the values from the summary.
+    ev = next(rr)
+    self.assertProtoEquals("""
+      value { tag: 'c1' simple_value: 1.0 }
+      value { tag: 'c2' simple_value: 2.0 }
+      value { tag: 'c3' simple_value: 3.0 }
+      """, ev.summary)
+
+    # We should be done.
+    self.assertRaises(StopIteration, lambda: next(rr))
+
+  def _csv_data(self, logdir):
+    # Create a small data file with 3 CSV records.
+    data_path = os.path.join(logdir, "data.csv")
+    with open(data_path, "w") as f:
+      f.write("1,2,3\n")
+      f.write("4,5,6\n")
+      f.write("7,8,9\n")
+    return data_path
+
+  def testManagedEndOfInputOneQueue(self):
+    # Tests that the supervisor finishes without an error when using
+    # a fixed number of epochs, reading from a single queue.
+    logdir = self._test_dir("managed_end_of_input_one_queue")
+    os.makedirs(logdir)
+    data_path = self._csv_data(logdir)
+    with ops.Graph().as_default():
+      # Create an input pipeline that reads the file 3 times.
+      filename_queue = input_lib.string_input_producer(
+          [data_path], num_epochs=3)
+      reader = io_ops.TextLineReader()
+      _, csv = reader.read(filename_queue)
+      rec = parsing_ops.decode_csv(csv, record_defaults=[[1], [1], [1]])
+      sv = supervisor.Supervisor(logdir=logdir)
+      with sv.managed_session("") as sess:
+        while not sv.should_stop():
+          sess.run(rec)
+
+  def testManagedEndOfInputTwoQueues(self):
+    # Tests that the supervisor finishes without an error when using
+    # a fixed number of epochs, reading from two queues, the second
+    # one producing a batch from the first one.
+    logdir = self._test_dir("managed_end_of_input_two_queues")
+    os.makedirs(logdir)
+    data_path = self._csv_data(logdir)
+    with ops.Graph().as_default():
+      # Create an input pipeline that reads the file 3 times.
+      filename_queue = input_lib.string_input_producer(
+          [data_path], num_epochs=3)
+      reader = io_ops.TextLineReader()
+      _, csv = reader.read(filename_queue)
+      rec = parsing_ops.decode_csv(csv, record_defaults=[[1], [1], [1]])
+      shuff_rec = input_lib.shuffle_batch(rec, 1, 6, 4)
+      sv = supervisor.Supervisor(logdir=logdir)
+      with sv.managed_session("") as sess:
+        while not sv.should_stop():
+          sess.run(shuff_rec)
+
+  def testManagedMainErrorTwoQueues(self):
+    # Tests that the supervisor correctly raises a main loop
+    # error even when using multiple queues for input.
+    logdir = self._test_dir("managed_main_error_two_queues")
+    os.makedirs(logdir)
+    data_path = self._csv_data(logdir)
+    with self.assertRaisesRegex(RuntimeError, "fail at step 3"):
+      with ops.Graph().as_default():
+        # Create an input pipeline that reads the file 3 times.
+        filename_queue = input_lib.string_input_producer(
+            [data_path], num_epochs=3)
+        reader = io_ops.TextLineReader()
+        _, csv = reader.read(filename_queue)
+        rec = parsing_ops.decode_csv(csv, record_defaults=[[1], [1], [1]])
+        shuff_rec = input_lib.shuffle_batch(rec, 1, 6, 4)
+        sv = supervisor.Supervisor(logdir=logdir)
+        with sv.managed_session("") as sess:
+          for step in range(9):
+            if sv.should_stop():
+              break
+            elif step == 3:
+              raise RuntimeError("fail at step 3")
+            else:
+              sess.run(shuff_rec)
+
+  def testSessionConfig(self):
+    logdir = self._test_dir("session_config")
+    with ops.Graph().as_default():
+      with ops.device("/cpu:1"):
+        my_op = constant_op.constant([1.0])
+      sv = supervisor.Supervisor(logdir=logdir)
+      sess = sv.prepare_or_wait_for_session(
+          "", config=config_pb2.ConfigProto(device_count={"CPU": 2}))
+      for _ in range(10):
+        self.evaluate(my_op)
+      sess.close()
+      sv.stop()
+
+  def testChiefCanWriteEvents(self):
+    logdir = self._test_dir("can_write")
+    with ops.Graph().as_default():
+      summary.scalar("c1", constant_op.constant(1))
+      summary.scalar("c2", constant_op.constant(2))
+      summary.scalar("c3", constant_op.constant(3))
+      summ = summary.merge_all()
+      sv = supervisor.Supervisor(is_chief=True, logdir=logdir, summary_op=None)
+      meta_graph_def = meta_graph.create_meta_graph_def()
+      sess = sv.prepare_or_wait_for_session("")
+      sv.summary_computed(sess, sess.run(summ))
+      sess.close()
+      # Wait to make sure everything is written to file before stopping.
+      time.sleep(1)
+      sv.stop()
+
+    rr = _summary_iterator(logdir)
+
+    # The first event should list the file_version.
+    ev = next(rr)
+    self.assertEqual("brain.Event:2", ev.file_version)
+
+    # The next one has the graph.
+    ev = next(rr)
+    ev_graph = graph_pb2.GraphDef()
+    ev_graph.ParseFromString(ev.graph_def)
+    self.assertProtoEquals(sess.graph.as_graph_def(add_shapes=True), ev_graph)
+
+    # Stored MetaGraphDef
+    ev = next(rr)
+    ev_meta_graph = meta_graph_pb2.MetaGraphDef()
+    ev_meta_graph.ParseFromString(ev.meta_graph_def)
+    self.assertProtoEquals(meta_graph_def, ev_meta_graph)
+    self.assertProtoEquals(
+        sess.graph.as_graph_def(add_shapes=True), ev_meta_graph.graph_def)
+    # The next one should have the values from the summary.
+    ev = next(rr)
+    self.assertProtoEquals("""
+      value { tag: 'c1' simple_value: 1.0 }
+      value { tag: 'c2' simple_value: 2.0 }
+      value { tag: 'c3' simple_value: 3.0 }
+      """, ev.summary)
+
+    # The next one should be a stop message if we closed cleanly.
+    ev = next(rr)
+    self.assertEqual(event_pb2.SessionLog.STOP, ev.session_log.status)
+
+    # We should be done.
+    self.assertRaises(StopIteration, lambda: next(rr))
+
+  def testNonChiefCannotWriteEvents(self):
+
+    def _summary_computed():
+      with ops.Graph().as_default():
+        sv = supervisor.Supervisor(is_chief=False)
+        sess = sv.prepare_or_wait_for_session("")
+        summary.scalar("c1", constant_op.constant(1))
+        summary.scalar("c2", constant_op.constant(2))
+        summ = summary.merge_all()
+        sv.summary_computed(sess, sess.run(summ))
+
+    def _start_standard_services():
+      with ops.Graph().as_default():
+        sv = supervisor.Supervisor(is_chief=False)
+        sess = sv.prepare_or_wait_for_session("")
+        sv.start_standard_services(sess)
+
+    self.assertRaises(RuntimeError, _summary_computed)
+    self.assertRaises(RuntimeError, _start_standard_services)
+
+  def testNoLogdirButWantSummary(self):
+    with ops.Graph().as_default():
+      summary.scalar("c1", constant_op.constant(1))
+      summary.scalar("c2", constant_op.constant(2))
+      summary.scalar("c3", constant_op.constant(3))
+      summ = summary.merge_all()
+      sv = supervisor.Supervisor(logdir="", summary_op=None)
+      sess = sv.prepare_or_wait_for_session("")
+      with self.assertRaisesRegex(RuntimeError, "requires a summary writer"):
+        sv.summary_computed(sess, sess.run(summ))
+
+  @test_util.run_v1_only("train.Supervisor is for v1 only")
+  def testLogdirButExplicitlyNoSummaryWriter(self):
+    logdir = self._test_dir("explicit_no_summary_writer")
+    with ops.Graph().as_default():
+      variable_v1.VariableV1([1.0], name="foo")
+      summary.scalar("c1", constant_op.constant(1))
+      summary.scalar("c2", constant_op.constant(2))
+      summary.scalar("c3", constant_op.constant(3))
+      summ = summary.merge_all()
+      sv = supervisor.Supervisor(logdir=logdir, summary_writer=None)
+      sess = sv.prepare_or_wait_for_session("")
+      # Check that a checkpoint is still be generated.
+      self._wait_for_glob(sv.save_path, 3.0)
+      # Check that we cannot write a summary
+      with self.assertRaisesRegex(RuntimeError, "requires a summary writer"):
+        sv.summary_computed(sess, sess.run(summ))
+
+  def testNoLogdirButExplicitSummaryWriter(self):
+    logdir = self._test_dir("explicit_summary_writer")
+    with ops.Graph().as_default():
+      summary.scalar("c1", constant_op.constant(1))
+      summary.scalar("c2", constant_op.constant(2))
+      summary.scalar("c3", constant_op.constant(3))
+      summ = summary.merge_all()
+      sw = writer.FileWriter(logdir)
+      sv = supervisor.Supervisor(logdir="", summary_op=None, summary_writer=sw)
+      meta_graph_def = meta_graph.create_meta_graph_def()
+      sess = sv.prepare_or_wait_for_session("")
+      sv.summary_computed(sess, sess.run(summ))
+      sess.close()
+      # Wait to make sure everything is written to file before stopping.
+      time.sleep(1)
+      sv.stop()
+
+    # Check the summary was written to 'logdir'
+    rr = _summary_iterator(logdir)
+
+    # The first event should list the file_version.
+    ev = next(rr)
+    self.assertEqual("brain.Event:2", ev.file_version)
+
+    # The next one has the graph.
+    ev = next(rr)
+    ev_graph = graph_pb2.GraphDef()
+    ev_graph.ParseFromString(ev.graph_def)
+    self.assertProtoEquals(sess.graph.as_graph_def(add_shapes=True), ev_graph)
+
+    # Stored MetaGraphDef
+    ev = next(rr)
+    ev_meta_graph = meta_graph_pb2.MetaGraphDef()
+    ev_meta_graph.ParseFromString(ev.meta_graph_def)
+    self.assertProtoEquals(meta_graph_def, ev_meta_graph)
+    self.assertProtoEquals(
+        sess.graph.as_graph_def(add_shapes=True), ev_meta_graph.graph_def)
+
+    # The next one should have the values from the summary.
+    ev = next(rr)
+    self.assertProtoEquals("""
+      value { tag: 'c1' simple_value: 1.0 }
+      value { tag: 'c2' simple_value: 2.0 }
+      value { tag: 'c3' simple_value: 3.0 }
+      """, ev.summary)
+
+    # The next one should be a stop message if we closed cleanly.
+    ev = next(rr)
+    self.assertEqual(event_pb2.SessionLog.STOP, ev.session_log.status)
+
+    # We should be done.
+    self.assertRaises(StopIteration, lambda: next(rr))
+
+  def testNoLogdirSucceeds(self):
+    with ops.Graph().as_default():
+      variable_v1.VariableV1([1.0, 2.0, 3.0])
+      sv = supervisor.Supervisor(logdir="", summary_op=None)
+      sess = sv.prepare_or_wait_for_session("")
+      sess.close()
+      sv.stop()
+
+  def testUseSessionManager(self):
+    with ops.Graph().as_default():
+      variable_v1.VariableV1([1.0, 2.0, 3.0])
+      sm = session_manager_lib.SessionManager()
+      # Pass in session_manager. The additional init_op is ignored.
+      sv = supervisor.Supervisor(logdir="", session_manager=sm)
+      sv.prepare_or_wait_for_session("")
+
+  @test_util.run_v1_only("train.Supervisor is for v1 only")
+  def testInitOp(self):
+    logdir = self._test_dir("default_init_op")
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1([1.0, 2.0, 3.0])
+      sv = supervisor.Supervisor(logdir=logdir)
+      sess = sv.prepare_or_wait_for_session("")
+      self.assertAllClose([1.0, 2.0, 3.0], sess.run(v))
+      sv.stop()
+
+  @test_util.run_v1_only("train.Supervisor is for v1 only")
+  def testInitFn(self):
+    logdir = self._test_dir("default_init_op")
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1([1.0, 2.0, 3.0])
+
+      def _init_fn(sess):
+        sess.run(v.initializer)
+
+      sv = supervisor.Supervisor(logdir=logdir, init_op=None, init_fn=_init_fn)
+      sess = sv.prepare_or_wait_for_session("")
+      self.assertAllClose([1.0, 2.0, 3.0], sess.run(v))
+      sv.stop()
+
+  @test_util.run_v1_only("train.Supervisor is for v1 only")
+  def testInitOpWithFeedDict(self):
+    logdir = self._test_dir("feed_dict_init_op")
+    with ops.Graph().as_default():
+      p = array_ops.placeholder(dtypes.float32, shape=(3,))
+      v = variable_v1.VariableV1(p, name="v")
+      sv = supervisor.Supervisor(
+          logdir=logdir,
+          init_op=variables.global_variables_initializer(),
+          init_feed_dict={p: [1.0, 2.0, 3.0]})
+      sess = sv.prepare_or_wait_for_session("")
+      self.assertAllClose([1.0, 2.0, 3.0], sess.run(v))
+      sv.stop()
+
+  @test_util.run_v1_only("train.Supervisor is for v1 only")
+  def testReadyForLocalInitOp(self):
+    server = server_lib.Server.create_local_server()
+    logdir = self._test_dir("default_ready_for_local_init_op")
+
+    uid = uuid.uuid4().hex
+
+    def get_session(is_chief):
+      g = ops.Graph()
+      with g.as_default():
+        with ops.device("/job:localhost"):
+          v = variable_v1.VariableV1(
+              1, name="default_ready_for_local_init_op_v_" + str(uid))
+          vadd = v.assign_add(1)
+          w = variable_v1.VariableV1(
+              v,
+              trainable=False,
+              collections=[ops.GraphKeys.LOCAL_VARIABLES],
+              name="default_ready_for_local_init_op_w_" + str(uid))
+          ready_for_local_init_op = variables.report_uninitialized_variables(
+              variables.global_variables())
+      sv = supervisor.Supervisor(
+          logdir=logdir,
+          is_chief=is_chief,
+          graph=g,
+          recovery_wait_secs=1,
+          init_op=v.initializer,
+          ready_for_local_init_op=ready_for_local_init_op)
+      sess = sv.prepare_or_wait_for_session(server.target)
+
+      return sv, sess, v, vadd, w
+
+    sv0, sess0, v0, _, w0 = get_session(True)
+    sv1, sess1, _, vadd1, w1 = get_session(False)
+
+    self.assertEqual(1, sess0.run(w0))
+    self.assertEqual(2, sess1.run(vadd1))
+    self.assertEqual(1, sess1.run(w1))
+    self.assertEqual(2, sess0.run(v0))
+
+    sv0.stop()
+    sv1.stop()
+
+  @test_util.run_v1_only("train.Supervisor is for v1 only")
+  def testReadyForLocalInitOpRestoreFromCheckpoint(self):
+    server = server_lib.Server.create_local_server()
+    logdir = self._test_dir("ready_for_local_init_op_restore")
+
+    uid = uuid.uuid4().hex
+
+    # Create a checkpoint.
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1(
+          10.0, name="ready_for_local_init_op_restore_v_" + str(uid))
+      summary.scalar("ready_for_local_init_op_restore_v_" + str(uid), v)
+      sv = supervisor.Supervisor(logdir=logdir)
+      sv.prepare_or_wait_for_session(server.target)
+      save_path = sv.save_path
+      self._wait_for_glob(save_path, 3.0)
+      self._wait_for_glob(
+          os.path.join(logdir, "*events*"), 3.0, for_checkpoint=False)
+      # Wait to make sure everything is written to file before stopping.
+      time.sleep(1)
+      sv.stop()
+
+    def get_session(is_chief):
+      g = ops.Graph()
+      with g.as_default():
+        with ops.device("/job:localhost"):
+          v = variable_v1.VariableV1(
+              1.0, name="ready_for_local_init_op_restore_v_" + str(uid))
+          vadd = v.assign_add(1)
+          w = variable_v1.VariableV1(
+              v,
+              trainable=False,
+              collections=[ops.GraphKeys.LOCAL_VARIABLES],
+              name="ready_for_local_init_op_restore_w_" + str(uid))
+          ready_for_local_init_op = variables.report_uninitialized_variables(
+              variables.global_variables())
+      sv = supervisor.Supervisor(
+          logdir=logdir,
+          is_chief=is_chief,
+          graph=g,
+          recovery_wait_secs=1,
+          ready_for_local_init_op=ready_for_local_init_op)
+      sess = sv.prepare_or_wait_for_session(server.target)
+
+      return sv, sess, v, vadd, w
+
+    sv0, sess0, v0, _, w0 = get_session(True)
+    sv1, sess1, _, vadd1, w1 = get_session(False)
+
+    self.assertEqual(10, sess0.run(w0))
+    self.assertEqual(11, sess1.run(vadd1))
+    self.assertEqual(10, sess1.run(w1))
+    self.assertEqual(11, sess0.run(v0))
+
+    sv0.stop()
+    sv1.stop()
+
+  def testLocalInitOp(self):
+    logdir = self._test_dir("default_local_init_op")
+    with ops.Graph().as_default():
+      # A local variable.
+      v = variable_v1.VariableV1([1.0, 2.0, 3.0],
+                                 trainable=False,
+                                 collections=[ops.GraphKeys.LOCAL_VARIABLES])
+
+      # An entity which is initialized through a TABLE_INITIALIZER.
+      w = variable_v1.VariableV1([4, 5, 6], trainable=False, collections=[])
+      ops.add_to_collection(ops.GraphKeys.TABLE_INITIALIZERS, w.initializer)
+
+      # This shouldn't add a variable to the VARIABLES collection responsible
+      # for variables that are saved/restored from checkpoints.
+      self.assertEqual(len(variables.global_variables()), 0)
+
+      # Suppress normal variable inits to make sure the local one is
+      # initialized via local_init_op.
+      sv = supervisor.Supervisor(logdir=logdir, init_op=None)
+      sess = sv.prepare_or_wait_for_session("")
+      self.assertAllClose([1.0, 2.0, 3.0], sess.run(v))
+      self.assertAllClose([4, 5, 6], sess.run(w))
+      sv.stop()
+
+  def testLocalInitOpForNonChief(self):
+    logdir = self._test_dir("default_local_init_op_non_chief")
+    with ops.Graph().as_default():
+      with ops.device("/job:localhost"):
+        # A local variable.
+        v = variable_v1.VariableV1([1.0, 2.0, 3.0],
+                                   trainable=False,
+                                   collections=[ops.GraphKeys.LOCAL_VARIABLES])
+        # This shouldn't add a variable to the VARIABLES collection responsible
+        # for variables that are saved/restored from checkpoints.
+        self.assertEqual(len(variables.global_variables()), 0)
+
+      # Suppress normal variable inits to make sure the local one is
+      # initialized via local_init_op.
+      sv = supervisor.Supervisor(logdir=logdir, init_op=None, is_chief=False)
+      sess = sv.prepare_or_wait_for_session("")
+      self.assertAllClose([1.0, 2.0, 3.0], sess.run(v))
+      sv.stop()
+
+  def testInitOpFails(self):
+    server = server_lib.Server.create_local_server()
+    logdir = self._test_dir("default_init_op_fails")
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1([1.0, 2.0, 3.0], name="v")
+      variable_v1.VariableV1([4.0, 5.0, 6.0], name="w")
+      # w will not be initialized.
+      sv = supervisor.Supervisor(logdir=logdir, init_op=v.initializer)
+      with self.assertRaisesRegex(RuntimeError, "Variables not initialized: w"):
+        sv.prepare_or_wait_for_session(server.target)
+
+  def testInitOpFailsForTransientVariable(self):
+    server = server_lib.Server.create_local_server()
+    logdir = self._test_dir("default_init_op_fails_for_local_variable")
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1([1.0, 2.0, 3.0],
+                                 name="v",
+                                 collections=[ops.GraphKeys.LOCAL_VARIABLES])
+      variable_v1.VariableV1([1.0, 2.0, 3.0],
+                             name="w",
+                             collections=[ops.GraphKeys.LOCAL_VARIABLES])
+      # w will not be initialized.
+      sv = supervisor.Supervisor(logdir=logdir, local_init_op=v.initializer)
+      with self.assertRaisesRegex(RuntimeError, "Variables not initialized: w"):
+        sv.prepare_or_wait_for_session(server.target)
+
+  @test_util.run_v1_only("train.Supervisor is for v1 only")
+  def testSetupFail(self):
+    logdir = self._test_dir("setup_fail")
+    with ops.Graph().as_default():
+      variable_v1.VariableV1([1.0, 2.0, 3.0], name="v")
+      with self.assertRaisesRegex(ValueError, "must have their device set"):
+        supervisor.Supervisor(logdir=logdir, is_chief=False)
+    with ops.Graph().as_default(), ops.device("/job:ps"):
+      variable_v1.VariableV1([1.0, 2.0, 3.0], name="v")
+      supervisor.Supervisor(logdir=logdir, is_chief=False)
+
+  @test_util.run_v1_only("train.Supervisor is for v1 only")
+  def testDefaultGlobalStep(self):
+    logdir = self._test_dir("default_global_step")
+    with ops.Graph().as_default():
+      variable_v1.VariableV1(287, name="global_step")
+      sv = supervisor.Supervisor(logdir=logdir)
+      sess = sv.prepare_or_wait_for_session("")
+      self.assertEqual(287, sess.run(sv.global_step))
+      sv.stop()
+
+  @test_util.run_v1_only("train.Supervisor is for v1 only")
+  def testRestoreFromMetaGraph(self):
+    logdir = self._test_dir("restore_from_meta_graph")
+    with ops.Graph().as_default():
+      variable_v1.VariableV1(1, name="v0")
+      sv = supervisor.Supervisor(logdir=logdir)
+      sess = sv.prepare_or_wait_for_session("")
+      filename = sv.saver.save(sess, sv.save_path)
+      sv.stop()
+    # Create a new Graph and Supervisor and recover.
+    with ops.Graph().as_default():
+      new_saver = saver_lib.import_meta_graph(".".join([filename, "meta"]))
+      self.assertIsNotNone(new_saver)
+      sv2 = supervisor.Supervisor(logdir=logdir, saver=new_saver)
+      sess = sv2.prepare_or_wait_for_session("")
+      self.assertEqual(1, sess.run("v0:0"))
+      sv2.saver.save(sess, sv2.save_path)
+      sv2.stop()
+
+  # This test is based on the fact that the standard services start
+  # right away and get to run once before sv.stop() returns.
+  # We still sleep a bit to make the test robust.
+  @test_util.run_v1_only("train.Supervisor is for v1 only")
+  def testStandardServicesWithoutGlobalStep(self):
+    logdir = self._test_dir("standard_services_without_global_step")
+    # Create a checkpoint.
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1([1.0], name="foo")
+      summary.scalar("v", v[0])
+      sv = supervisor.Supervisor(logdir=logdir)
+      meta_graph_def = meta_graph.create_meta_graph_def(
+          saver_def=sv.saver.saver_def)
+      sess = sv.prepare_or_wait_for_session("")
+      save_path = sv.save_path
+      self._wait_for_glob(save_path, 3.0)
+      self._wait_for_glob(
+          os.path.join(logdir, "*events*"), 3.0, for_checkpoint=False)
+      # Wait to make sure everything is written to file before stopping.
+      time.sleep(1)
+      sv.stop()
+    # There should be an event file with a version number.
+    rr = _summary_iterator(logdir)
+    ev = next(rr)
+    self.assertEqual("brain.Event:2", ev.file_version)
+    ev = next(rr)
+    ev_graph = graph_pb2.GraphDef()
+    ev_graph.ParseFromString(ev.graph_def)
+    self.assertProtoEquals(sess.graph.as_graph_def(add_shapes=True), ev_graph)
+
+    # Stored MetaGraphDef
+    ev = next(rr)
+    ev_meta_graph = meta_graph_pb2.MetaGraphDef()
+    ev_meta_graph.ParseFromString(ev.meta_graph_def)
+    self.assertProtoEquals(meta_graph_def, ev_meta_graph)
+    self.assertProtoEquals(
+        sess.graph.as_graph_def(add_shapes=True), ev_meta_graph.graph_def)
+
+    ev = next(rr)
+    self.assertProtoEquals("value { tag: 'v' simple_value: 1.0 }", ev.summary)
+
+    ev = next(rr)
+    self.assertEqual(event_pb2.SessionLog.STOP, ev.session_log.status)
+
+    self.assertRaises(StopIteration, lambda: next(rr))
+    # There should be a checkpoint file with the variable "foo"
+    with ops.Graph().as_default(), self.cached_session() as sess:
+      v = variable_v1.VariableV1([10.10], name="foo")
+      sav = saver_lib.Saver([v])
+      sav.restore(sess, save_path)
+      self.assertEqual(1.0, self.evaluate(v)[0])
+
+  # Same as testStandardServicesNoGlobalStep but with a global step.
+  # We should get a summary about the step time.
+  @test_util.run_v1_only("train.Supervisor is for v1 only")
+  def testStandardServicesWithGlobalStep(self):
+    logdir = self._test_dir("standard_services_with_global_step")
+    # Create a checkpoint.
+    with ops.Graph().as_default():
+      v = variable_v1.VariableV1([123], name="global_step")
+      sv = supervisor.Supervisor(logdir=logdir)
+      meta_graph_def = meta_graph.create_meta_graph_def(
+          saver_def=sv.saver.saver_def)
+      sess = sv.prepare_or_wait_for_session("")
+      # This is where the checkpoint will appear, with step number 123.
+      save_path = "%s-123" % sv.save_path
+      self._wait_for_glob(save_path, 3.0)
+      self._wait_for_glob(
+          os.path.join(logdir, "*events*"), 3.0, for_checkpoint=False)
+      # Wait to make sure everything is written to file before stopping.
+      time.sleep(1)
+      sv.stop()
+    # There should be an event file with a version number.
+    rr = _summary_iterator(logdir)
+    ev = next(rr)
+    self.assertEqual("brain.Event:2", ev.file_version)
+    ev = next(rr)
+    ev_graph = graph_pb2.GraphDef()
+    ev_graph.ParseFromString(ev.graph_def)
+    self.assertProtoEquals(sess.graph.as_graph_def(add_shapes=True), ev_graph)
+    ev = next(rr)
+    ev_meta_graph = meta_graph_pb2.MetaGraphDef()
+    ev_meta_graph.ParseFromString(ev.meta_graph_def)
+    self.assertProtoEquals(meta_graph_def, ev_meta_graph)
+    self.assertProtoEquals(
+        sess.graph.as_graph_def(add_shapes=True), ev_meta_graph.graph_def)
+    ev = next(rr)
+    # It is actually undeterministic whether SessionLog.START gets written
+    # before the summary or the checkpoint, but this works when run 10000 times.
+    self.assertEqual(123, ev.step)
+    self.assertEqual(event_pb2.SessionLog.START, ev.session_log.status)
+    first = next(rr)
+    second = next(rr)
+    # It is undeterministic whether the value gets written before the checkpoint
+    # since they are on separate threads, so we check for both conditions.
+    if first.HasField("summary"):
+      self.assertProtoEquals("""value { tag: 'global_step/sec'
+                                        simple_value: 0.0 }""", first.summary)
+      self.assertEqual(123, second.step)
+      self.assertEqual(event_pb2.SessionLog.CHECKPOINT,
+                       second.session_log.status)
+    else:
+      self.assertEqual(123, first.step)
+      self.assertEqual(event_pb2.SessionLog.CHECKPOINT,
+                       first.session_log.status)
+      self.assertProtoEquals("""value { tag: 'global_step/sec'
+                                        simple_value: 0.0 }""", second.summary)
+    ev = next(rr)
+    self.assertEqual(event_pb2.SessionLog.STOP, ev.session_log.status)
+    self.assertRaises(StopIteration, lambda: next(rr))
+    # There should be a checkpoint file with the variable "foo"
+    with ops.Graph().as_default(), self.cached_session() as sess:
+      v = variable_v1.VariableV1([-12], name="global_step")
+      sav = saver_lib.Saver([v])
+      sav.restore(sess, save_path)
+      self.assertEqual(123, self.evaluate(v)[0])
+
+  def testNoQueueRunners(self):
+    with ops.Graph().as_default(), self.cached_session() as sess:
+      sv = supervisor.Supervisor(logdir=self._test_dir("no_queue_runners"))
+      self.assertEqual(0, len(sv.start_queue_runners(sess)))
+      sv.stop()
+
+  def testPrepareSessionAfterStopForChief(self):
+    logdir = self._test_dir("prepare_after_stop_chief")
+    with ops.Graph().as_default():
+      sv = supervisor.Supervisor(logdir=logdir, is_chief=True)
+
+      # Create a first session and then stop.
+      sess = sv.prepare_or_wait_for_session("")
+      sv.stop()
+      sess.close()
+      self.assertTrue(sv.should_stop())
+
+      # Now create a second session and test that we don't stay stopped, until
+      # we ask to stop again.
+      sess2 = sv.prepare_or_wait_for_session("")
+      self.assertFalse(sv.should_stop())
+      sv.stop()
+      sess2.close()
+      self.assertTrue(sv.should_stop())
+
+  def testPrepareSessionAfterStopForNonChief(self):
+    logdir = self._test_dir("prepare_after_stop_nonchief")
+    with ops.Graph().as_default():
+      sv = supervisor.Supervisor(logdir=logdir, is_chief=False)
+
+      # Create a first session and then stop.
+      sess = sv.prepare_or_wait_for_session("")
+      sv.stop()
+      sess.close()
+      self.assertTrue(sv.should_stop())
+
+      # Now create a second session and test that we don't stay stopped, until
+      # we ask to stop again.
+      sess2 = sv.prepare_or_wait_for_session("")
+      self.assertFalse(sv.should_stop())
+      sv.stop()
+      sess2.close()
+      self.assertTrue(sv.should_stop())
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/sync_replicas_optimizer.py

@@ -0,0 +1,494 @@
+# Copyright 2016 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.
+# ==============================================================================
+
+"""Synchronize replicas for training."""
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import data_flow_ops
+from tensorflow.python.ops import state_ops
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.training import optimizer
+from tensorflow.python.training import queue_runner
+from tensorflow.python.training import session_manager
+from tensorflow.python.training import session_run_hook
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+# Please note that the gradients from replicas are averaged instead of summed
+# (as in the old sync_replicas_optimizer) so you need to increase the learning
+# rate according to the number of replicas. This change is introduced to be
+# consistent with how gradients are aggregated (averaged) within a batch in a
+# replica.
+@tf_export(v1=["train.SyncReplicasOptimizer"])
+class SyncReplicasOptimizer(optimizer.Optimizer):
+  """Class to synchronize, aggregate gradients and pass them to the optimizer.
+
+  This class is deprecated. For synchronous training, please use [Distribution
+  Strategies](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/distribute).
+
+  In a typical asynchronous training environment, it's common to have some
+  stale gradients. For example, with a N-replica asynchronous training,
+  gradients will be applied to the variables N times independently. Depending
+  on each replica's training speed, some gradients might be calculated from
+  copies of the variable from several steps back (N-1 steps on average). This
+  optimizer avoids stale gradients by collecting gradients from all replicas,
+  averaging them, then applying them to the variables in one shot, after
+  which replicas can fetch the new variables and continue.
+
+  The following accumulators/queue are created:
+
+  * N `gradient accumulators`, one per variable to train. Gradients are pushed
+    to them and the chief worker will wait until enough gradients are collected
+    and then average them before applying to variables. The accumulator will
+    drop all stale gradients (more details in the accumulator op).
+  * 1 `token` queue where the optimizer pushes the new global_step value after
+    all variables are updated.
+
+  The following local variable is created:
+  * `sync_rep_local_step`, one per replica. Compared against the global_step in
+    each accumulator to check for staleness of the gradients.
+
+  The optimizer adds nodes to the graph to collect gradients and pause the
+  trainers until variables are updated.
+  For the Parameter Server job:
+
+  1. An accumulator is created for each variable, and each replica pushes the
+     gradients into the accumulators instead of directly applying them to the
+     variables.
+  2. Each accumulator averages once enough gradients (replicas_to_aggregate)
+     have been accumulated.
+  3. Apply the averaged gradients to the variables.
+  4. Only after all variables have been updated, increment the global step.
+  5. Only after step 4, pushes `global_step` in the `token_queue`, once for
+     each worker replica. The workers can now fetch the global step, use it to
+     update its local_step variable and start the next batch. Please note that
+     some workers can consume multiple minibatches, while some may not consume
+     even one. This is because each worker fetches minibatches as long as
+     a token exists. If one worker is stuck for some reason and does not
+     consume a token, another worker can use it.
+
+  For the replicas:
+
+  1. Start a step: fetch variables and compute gradients.
+  2. Once the gradients have been computed, push them into gradient
+     accumulators. Each accumulator will check the staleness and drop the stale.
+  3. After pushing all the gradients, dequeue an updated value of global_step
+     from the token queue and record that step to its local_step variable. Note
+     that this is effectively a barrier.
+  4. Start the next batch.
+
+  ### Usage
+
+  ```python
+  # Create any optimizer to update the variables, say a simple SGD:
+  opt = GradientDescentOptimizer(learning_rate=0.1)
+
+  # Wrap the optimizer with sync_replicas_optimizer with 50 replicas: at each
+  # step the optimizer collects 50 gradients before applying to variables.
+  # Note that if you want to have 2 backup replicas, you can change
+  # total_num_replicas=52 and make sure this number matches how many physical
+  # replicas you started in your job.
+  opt = tf.compat.v1.train.SyncReplicasOptimizer(opt, replicas_to_aggregate=50,
+                                 total_num_replicas=50)
+
+  # Some models have startup_delays to help stabilize the model but when using
+  # sync_replicas training, set it to 0.
+
+  # Now you can call `minimize()` or `compute_gradients()` and
+  # `apply_gradients()` normally
+  training_op = opt.minimize(total_loss, global_step=self.global_step)
+
+
+  # You can create the hook which handles initialization and queues.
+  sync_replicas_hook = opt.make_session_run_hook(is_chief)
+  ```
+
+  In the training program, every worker will run the train_op as if not
+  synchronized.
+
+  ```python
+  with training.MonitoredTrainingSession(
+      master=workers[worker_id].target, is_chief=is_chief,
+      hooks=[sync_replicas_hook]) as mon_sess:
+    while not mon_sess.should_stop():
+      mon_sess.run(training_op)
+  ```
+  """
+
+  @deprecation.deprecated(
+      None, "The `SyncReplicaOptimizer` class is deprecated. For synchronous "
+      "training, please use [Distribution Strategies](https://github.com/"
+      "tensorflow/tensorflow/tree/master/tensorflow/contrib/distribute).",
+      warn_once=True)
+  def __init__(self,
+               opt,
+               replicas_to_aggregate,
+               total_num_replicas=None,
+               variable_averages=None,
+               variables_to_average=None,
+               use_locking=False,
+               name="sync_replicas"):
+    """Construct a sync_replicas optimizer.
+
+    Args:
+      opt: The actual optimizer that will be used to compute and apply the
+        gradients. Must be one of the Optimizer classes.
+      replicas_to_aggregate: number of replicas to aggregate for each variable
+        update.
+      total_num_replicas: Total number of tasks/workers/replicas, could be
+        different from replicas_to_aggregate.
+        If total_num_replicas > replicas_to_aggregate: it is backup_replicas +
+        replicas_to_aggregate.
+        If total_num_replicas < replicas_to_aggregate: Replicas compute
+        multiple batches per update to variables.
+      variable_averages: Optional `ExponentialMovingAverage` object, used to
+        maintain moving averages for the variables passed in
+        `variables_to_average`.
+      variables_to_average: a list of variables that need to be averaged. Only
+        needed if variable_averages is passed in.
+      use_locking: If True use locks for update operation.
+      name: string. Optional name of the returned operation.
+    """
+    if total_num_replicas is None:
+      total_num_replicas = replicas_to_aggregate
+
+    super(SyncReplicasOptimizer, self).__init__(use_locking, name)
+    logging.info(
+        "SyncReplicasV2: replicas_to_aggregate=%s; total_num_replicas=%s",
+        replicas_to_aggregate, total_num_replicas)
+    self._opt = opt
+    self._replicas_to_aggregate = replicas_to_aggregate
+    self._gradients_applied = False
+    self._variable_averages = variable_averages
+    self._variables_to_average = variables_to_average
+    self._total_num_replicas = total_num_replicas
+    self._tokens_per_step = max(total_num_replicas, replicas_to_aggregate)
+    self._global_step = None
+    self._sync_token_queue = None
+
+    # The synchronization op will be executed in a queue runner which should
+    # only be executed by one of the replicas (usually the chief).
+    self._chief_queue_runner = None
+
+    # Remember which accumulator is on which device to set the initial step in
+    # the accumulator to be global step. This list contains list of the
+    # following format: (accumulator, device).
+    self._accumulator_list = []
+
+  def compute_gradients(self, *args, **kwargs):
+    """Compute gradients of "loss" for the variables in "var_list".
+
+    This simply wraps the compute_gradients() from the real optimizer. The
+    gradients will be aggregated in the apply_gradients() so that user can
+    modify the gradients like clipping with per replica global norm if needed.
+    The global norm with aggregated gradients can be bad as one replica's huge
+    gradients can hurt the gradients from other replicas.
+
+    Args:
+      *args: Arguments for compute_gradients().
+      **kwargs: Keyword arguments for compute_gradients().
+
+    Returns:
+      A list of (gradient, variable) pairs.
+    """
+    return self._opt.compute_gradients(*args, **kwargs)
+
+  def apply_gradients(self, grads_and_vars, global_step=None, name=None):
+    """Apply gradients to variables.
+
+    This contains most of the synchronization implementation and also wraps the
+    apply_gradients() from the real optimizer.
+
+    Args:
+      grads_and_vars: List of (gradient, variable) pairs as returned by
+        compute_gradients().
+      global_step: Optional Variable to increment by one after the
+        variables have been updated.
+      name: Optional name for the returned operation.  Default to the
+        name passed to the Optimizer constructor.
+
+    Returns:
+      train_op: The op to dequeue a token so the replicas can exit this batch
+      and start the next one. This is executed by each replica.
+
+    Raises:
+      ValueError: If the grads_and_vars is empty.
+      ValueError: If global step is not provided, the staleness cannot be
+        checked.
+    """
+    if not grads_and_vars:
+      raise ValueError("Must supply at least one variable")
+
+    if global_step is None:
+      raise ValueError("Global step is required to check staleness")
+
+    self._global_step = global_step
+    train_ops = []
+    aggregated_grad = []
+    var_list = []
+
+    # local_anchor op will be placed on this worker task by default.
+    local_anchor = control_flow_ops.no_op()
+    # Colocating local_step variable prevents it being placed on the PS.
+    distribution_strategy = distribute_lib.get_strategy()
+    with distribution_strategy.extended.colocate_vars_with(local_anchor):
+      self._local_step = variable_v1.VariableV1(
+          initial_value=0,
+          trainable=False,
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
+          dtype=global_step.dtype.base_dtype,
+          name="sync_rep_local_step")
+
+    self.local_step_init_op = state_ops.assign(self._local_step, global_step)
+    chief_init_ops = [self.local_step_init_op]
+    self.ready_for_local_init_op = variables.report_uninitialized_variables(
+        variables.global_variables())
+
+    with ops.name_scope(None, self._name):
+      for grad, var in grads_and_vars:
+        var_list.append(var)
+        with ops.device(var.device):
+          # Dense gradients.
+          if grad is None:
+            aggregated_grad.append(None)  # pass-through.
+            continue
+          elif isinstance(grad, tensor.Tensor):
+            grad_accum = data_flow_ops.ConditionalAccumulator(
+                grad.dtype,
+                shape=var.get_shape(),
+                shared_name=var.name + "/grad_accum")
+            train_ops.append(grad_accum.apply_grad(
+                grad, local_step=self._local_step))
+            aggregated_grad.append(grad_accum.take_grad(
+                self._replicas_to_aggregate))
+          else:
+            if not isinstance(grad, indexed_slices.IndexedSlices):
+              raise ValueError("Unknown grad type!")
+            grad_accum = data_flow_ops.SparseConditionalAccumulator(
+                grad.dtype, shape=(), shared_name=var.name + "/grad_accum")
+            train_ops.append(grad_accum.apply_indexed_slices_grad(
+                grad, local_step=self._local_step))
+            aggregated_grad.append(grad_accum.take_indexed_slices_grad(
+                self._replicas_to_aggregate))
+
+          self._accumulator_list.append((grad_accum, var.device))
+
+      aggregated_grads_and_vars = zip(aggregated_grad, var_list)
+
+      # sync_op will be assigned to the same device as the global step.
+      with ops.device(global_step.device), ops.name_scope(""):
+        update_op = self._opt.apply_gradients(aggregated_grads_and_vars,
+                                              global_step)
+
+      # Create token queue.
+      with ops.device(global_step.device), ops.name_scope(""):
+        sync_token_queue = (
+            data_flow_ops.FIFOQueue(-1,
+                                    global_step.dtype.base_dtype,
+                                    shapes=(),
+                                    name="sync_token_q",
+                                    shared_name="sync_token_q"))
+        self._sync_token_queue = sync_token_queue
+
+      with ops.device(global_step.device), ops.name_scope(""):
+        # Replicas have to wait until they can get a token from the token queue.
+        with ops.control_dependencies(train_ops):
+          token = sync_token_queue.dequeue()
+        train_op = state_ops.assign(self._local_step, token)
+
+        with ops.control_dependencies([update_op]):
+          # Sync_op needs to insert tokens to the token queue at the end of the
+          # step so the replicas can fetch them to start the next step.
+          tokens = array_ops.fill([self._tokens_per_step], global_step)
+          sync_op = sync_token_queue.enqueue_many((tokens,))
+
+        if self._variable_averages is not None:
+          with ops.control_dependencies([sync_op]), ops.name_scope(""):
+            sync_op = self._variable_averages.apply(
+                self._variables_to_average)
+
+        self._chief_queue_runner = queue_runner.QueueRunner(
+            sync_token_queue, [sync_op])
+      for accum, dev in self._accumulator_list:
+        with ops.device(dev):
+          chief_init_ops.append(
+              accum.set_global_step(
+                  global_step, name="SetGlobalStep"))
+      self.chief_init_op = control_flow_ops.group(*(chief_init_ops))
+      self._gradients_applied = True
+      return train_op
+
+  def get_chief_queue_runner(self):
+    """Returns the QueueRunner for the chief to execute.
+
+    This includes the operations to synchronize replicas: aggregate gradients,
+    apply to variables, increment global step, insert tokens to token queue.
+
+    Note that this can only be called after calling apply_gradients() which
+    actually generates this queuerunner.
+
+    Returns:
+      A `QueueRunner` for chief to execute.
+
+    Raises:
+      ValueError: If this is called before apply_gradients().
+    """
+    if self._gradients_applied is False:
+      raise ValueError("Should be called after apply_gradients().")
+
+    return self._chief_queue_runner
+
+  def get_slot(self, *args, **kwargs):
+    """Return a slot named "name" created for "var" by the Optimizer.
+
+    This simply wraps the get_slot() from the actual optimizer.
+
+    Args:
+      *args: Arguments for get_slot().
+      **kwargs: Keyword arguments for get_slot().
+
+    Returns:
+      The `Variable` for the slot if it was created, `None` otherwise.
+    """
+    return self._opt.get_slot(*args, **kwargs)
+
+  def variables(self):
+    """Fetches a list of optimizer variables in the default graph.
+
+    This wraps `variables()` from the actual optimizer. It does not include
+    the `SyncReplicasOptimizer`'s local step.
+
+    Returns:
+      A list of variables.
+    """
+    return self._opt.variables()
+
+  def get_slot_names(self, *args, **kwargs):
+    """Return a list of the names of slots created by the `Optimizer`.
+
+    This simply wraps the get_slot_names() from the actual optimizer.
+
+    Args:
+      *args: Arguments for get_slot().
+      **kwargs: Keyword arguments for get_slot().
+
+    Returns:
+      A list of strings.
+    """
+    return self._opt.get_slot_names(*args, **kwargs)
+
+  def get_init_tokens_op(self, num_tokens=-1):
+    """Returns the op to fill the sync_token_queue with the tokens.
+
+    This is supposed to be executed in the beginning of the chief/sync thread
+    so that even if the total_num_replicas is less than replicas_to_aggregate,
+    the model can still proceed as the replicas can compute multiple steps per
+    variable update. Make sure:
+    `num_tokens >= replicas_to_aggregate - total_num_replicas`.
+
+    Args:
+      num_tokens: Number of tokens to add to the queue.
+
+    Returns:
+      An op for the chief/sync replica to fill the token queue.
+
+    Raises:
+      ValueError: If this is called before apply_gradients().
+      ValueError: If num_tokens are smaller than replicas_to_aggregate -
+        total_num_replicas.
+    """
+    if self._gradients_applied is False:
+      raise ValueError(
+          "get_init_tokens_op() should be called after apply_gradients().")
+
+    tokens_needed = self._replicas_to_aggregate - self._total_num_replicas
+    if num_tokens == -1:
+      num_tokens = self._replicas_to_aggregate
+    elif num_tokens < tokens_needed:
+      raise ValueError(
+          "Too few tokens to finish the first step: %d (given) vs %d (needed)" %
+          (num_tokens, tokens_needed))
+
+    if num_tokens > 0:
+      with ops.device(self._global_step.device), ops.name_scope(""):
+        tokens = array_ops.fill([num_tokens], self._global_step)
+        init_tokens = self._sync_token_queue.enqueue_many((tokens,))
+    else:
+      init_tokens = control_flow_ops.no_op(name="no_init_tokens")
+
+    return init_tokens
+
+  def make_session_run_hook(self, is_chief, num_tokens=-1):
+    """Creates a hook to handle SyncReplicasHook ops such as initialization."""
+    return _SyncReplicasOptimizerHook(self, is_chief, num_tokens)
+
+
+class _SyncReplicasOptimizerHook(session_run_hook.SessionRunHook):
+  """A SessionRunHook handles ops related to SyncReplicasOptimizer."""
+
+  def __init__(self, sync_optimizer, is_chief, num_tokens):
+    """Creates hook to handle SyncReplicasOptimizer initialization ops.
+
+    Args:
+      sync_optimizer: `SyncReplicasOptimizer` which this hook will initialize.
+      is_chief: `Bool`, whether is this a chief replica or not.
+      num_tokens: Number of tokens to add to the queue.
+    """
+    self._sync_optimizer = sync_optimizer
+    self._is_chief = is_chief
+    self._num_tokens = num_tokens
+
+  def begin(self):
+    if self._sync_optimizer._gradients_applied is False:  # pylint: disable=protected-access
+      raise ValueError(
+          "SyncReplicasOptimizer.apply_gradient should be called before using "
+          "the hook.")
+    if self._is_chief:
+      self._local_init_op = self._sync_optimizer.chief_init_op
+      self._ready_for_local_init_op = (
+          self._sync_optimizer.ready_for_local_init_op)
+      self._q_runner = self._sync_optimizer.get_chief_queue_runner()
+      self._init_tokens_op = self._sync_optimizer.get_init_tokens_op(
+          self._num_tokens)
+    else:
+      self._local_init_op = self._sync_optimizer.local_step_init_op
+      self._ready_for_local_init_op = (
+          self._sync_optimizer.ready_for_local_init_op)
+      self._q_runner = None
+      self._init_tokens_op = None
+
+  def after_create_session(self, session, coord):
+    """Runs SyncReplicasOptimizer initialization ops."""
+    local_init_success, msg = session_manager._ready(  # pylint: disable=protected-access
+        self._ready_for_local_init_op, session,
+        "Model is not ready for SyncReplicasOptimizer local init.")
+    if not local_init_success:
+      raise RuntimeError(
+          "Init operations did not make model ready for SyncReplicasOptimizer "
+          "local_init. Init op: %s, error: %s" %
+          (self._local_init_op.name, msg))
+    session.run(self._local_init_op)
+    if self._init_tokens_op is not None:
+      session.run(self._init_tokens_op)
+    if self._q_runner is not None:
+      self._q_runner.create_threads(
+          session, coord=coord, daemon=True, start=True)

tensorflow/python/training/sync_replicas_optimizer_test.py

@@ -0,0 +1,301 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Tests for sync_replicas_optimizer.py."""
+
+import time
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.framework.test_util import create_local_cluster
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.platform import test
+from tensorflow.python.training import adam
+from tensorflow.python.training import gradient_descent
+from tensorflow.python.training import training
+
+
+# Creates the workers and return their sessions, graphs, train_ops.
+def get_workers(num_workers, replicas_to_aggregate, workers):
+  sessions = []
+  graphs = []
+  train_ops = []
+  for worker_id in range(num_workers):
+    graph = ops.Graph()
+    is_chief = (worker_id == 0)
+    with graph.as_default():
+      with ops.device("/job:ps/task:0"):
+        global_step = variable_v1.VariableV1(
+            0, name="global_step", trainable=False)
+        var_0 = variable_v1.VariableV1(0.0, name="v0")
+      with ops.device("/job:ps/task:1"):
+        var_1 = variable_v1.VariableV1(1.0, name="v1")
+        var_sparse = variable_v1.VariableV1([[3.0], [4.0]], name="v_sparse")
+
+      with ops.device("/job:worker/task:" + str(worker_id)):
+        grads_0 = constant_op.constant(0.1 + worker_id * 0.2)
+        grads_1 = constant_op.constant(0.9 + worker_id * 0.2)
+        # This is to test against sparse gradients.
+        grads_sparse = indexed_slices.IndexedSlices(
+            constant_op.constant(
+                [0.1 + worker_id * 0.2], shape=[1, 1]),
+            constant_op.constant([1]),
+            constant_op.constant([2, 1]))
+        sgd_opt = gradient_descent.GradientDescentOptimizer(2.0)
+        sync_rep_opt = training.SyncReplicasOptimizer(
+            sgd_opt,
+            replicas_to_aggregate=replicas_to_aggregate,
+            total_num_replicas=num_workers)
+        train_op = [
+            sync_rep_opt.apply_gradients(
+                zip([grads_0, grads_1, grads_sparse],
+                    [var_0, var_1, var_sparse]),
+                global_step=global_step)
+        ]
+        sync_replicas_hook = sync_rep_opt.make_session_run_hook(
+            is_chief, num_tokens=num_workers)
+
+      # Creates MonitoredSession
+      session = training.MonitoredTrainingSession(
+          master=workers[worker_id].target,
+          is_chief=is_chief,
+          hooks=[sync_replicas_hook])
+
+    sessions.append(session)
+    graphs.append(graph)
+    train_ops.append(train_op)
+
+  return sessions, graphs, train_ops
+
+
+class SyncReplicasOptimizerTest(test.TestCase):
+
+  def _run(self, train_op, sess):
+    sess.run(train_op)
+
+  @test_util.run_v1_only(
+      "This exercises tensor lookup via names which is not supported in V2.")
+  def test2Workers(self):
+    num_workers = 2
+    replicas_to_aggregate = 2
+    num_ps = 2
+    workers, _ = create_local_cluster(num_workers=num_workers, num_ps=num_ps)
+
+    # Creates and returns all the workers.
+    sessions, graphs, train_ops = get_workers(num_workers,
+                                              replicas_to_aggregate, workers)
+
+    # Chief should have already initialized all the variables.
+    var_0_g_0 = graphs[0].get_tensor_by_name("v0:0")
+    var_1_g_0 = graphs[0].get_tensor_by_name("v1:0")
+    local_step_0 = graphs[0].get_tensor_by_name("sync_rep_local_step:0")
+    self.assertAllEqual(0.0, sessions[0].run(var_0_g_0))
+    self.assertAllEqual(1.0, sessions[0].run(var_1_g_0))
+    self.assertAllEqual(0, sessions[0].run(local_step_0))
+
+    # Will just use session 1 to verify all the variables later.
+    var_0_g_1 = graphs[1].get_tensor_by_name("v0:0")
+    var_1_g_1 = graphs[1].get_tensor_by_name("v1:0")
+    var_sparse_g_1 = graphs[1].get_tensor_by_name("v_sparse:0")
+    local_step_1 = graphs[1].get_tensor_by_name("sync_rep_local_step:0")
+    global_step = graphs[1].get_tensor_by_name("global_step:0")
+
+    # The steps should also be initialized.
+    self.assertAllEqual(0, sessions[1].run(global_step))
+    self.assertAllEqual(0, sessions[1].run(local_step_1))
+    self.assertAllClose([[3.0], [4.0]], sessions[1].run(var_sparse_g_1))
+
+    # We have initial tokens in the queue so we can call this one by one. After
+    # the first step, this will no longer work as there will be no more extra
+    # tokens in the queue.
+    sessions[0].run(train_ops[0])
+    sessions[1].run(train_ops[1])
+
+    # The global step should have been updated and the variables should now have
+    # the new values after the average of the gradients are applied.
+    while sessions[1].run(global_step) != 1:
+      time.sleep(0.01)
+
+    self.assertAllClose(0 - (0.1 + 0.3) / 2 * 2.0, sessions[1].run(var_0_g_1))
+    self.assertAllClose(1 - (0.9 + 1.1) / 2 * 2.0, sessions[1].run(var_1_g_1))
+    self.assertAllClose([[3.0], [4.0 - (0.1 + 0.3) / 2 * 2.0]],
+                        sessions[1].run(var_sparse_g_1))
+
+    # The local step for both workers should still be 0 because the initial
+    # tokens in the token queue are 0s. This means that the following
+    # computation of the gradients will be wasted as local_step is smaller than
+    # the current global step. However, this only happens once when the system
+    # just starts and this is necessary to make the system robust for the case
+    # when chief gets restarted by errors/preemption/...
+    self.assertAllEqual(0, sessions[0].run(local_step_0))
+    self.assertAllEqual(0, sessions[1].run(local_step_1))
+
+    sessions[0].run(train_ops[0])
+    sessions[1].run(train_ops[1])
+    # Although the global step should still be 1 as explained above, the local
+    # step should now be updated to 1. The variables are still the same.
+    self.assertAllEqual(1, sessions[1].run(global_step))
+    self.assertAllEqual(1, sessions[0].run(local_step_0))
+    self.assertAllEqual(1, sessions[1].run(local_step_1))
+    self.assertAllClose(0 - (0.1 + 0.3) / 2 * 2.0, sessions[1].run(var_0_g_1))
+    self.assertAllClose(1 - (0.9 + 1.1) / 2 * 2.0, sessions[1].run(var_1_g_1))
+
+    # At this step, the token queue is empty. So the 2 workers need to work
+    # together to proceed.
+    threads = []
+    threads.append(
+        self.checkedThread(
+            target=self._run, args=(train_ops[0], sessions[0])))
+    threads.append(
+        self.checkedThread(
+            target=self._run, args=(train_ops[1], sessions[1])))
+
+    # The two workers starts to execute the train op.
+    for thread in threads:
+      thread.start()
+    for thread in threads:
+      thread.join()
+
+    # The global step should now be 2 and the gradients should have been
+    # applied twice.
+    self.assertAllEqual(2, sessions[1].run(global_step))
+    self.assertAllClose(0 - 2 * (0.1 + 0.3) / 2 * 2.0,
+                        sessions[1].run(var_0_g_1))
+    self.assertAllClose(1 - 2 * (0.9 + 1.1) / 2 * 2.0,
+                        sessions[1].run(var_1_g_1))
+
+  # 3 workers and one of them is backup.
+  @test_util.run_v1_only(
+      "This exercises tensor lookup via names which is not supported in V2.")
+  def test3Workers1Backup(self):
+    num_workers = 3
+    replicas_to_aggregate = 2
+    num_ps = 2
+    workers, _ = create_local_cluster(num_workers=num_workers, num_ps=num_ps)
+
+    # Creates and returns all the workers.
+    sessions, graphs, train_ops = get_workers(num_workers,
+                                              replicas_to_aggregate, workers)
+
+    # Chief should have already initialized all the variables.
+    var_0_g_1 = graphs[1].get_tensor_by_name("v0:0")
+    var_1_g_1 = graphs[1].get_tensor_by_name("v1:0")
+    local_step_1 = graphs[1].get_tensor_by_name("sync_rep_local_step:0")
+    global_step = graphs[1].get_tensor_by_name("global_step:0")
+
+    # The steps should also be initialized.
+    self.assertAllEqual(0, sessions[1].run(global_step))
+    self.assertAllEqual(0, sessions[1].run(local_step_1))
+
+    # We have initial tokens in the queue so we can call this one by one. After
+    # the token queue becomes empty, they should be called concurrently.
+    # Here worker 0 and worker 2 finished first.
+    sessions[0].run(train_ops[0])
+    sessions[2].run(train_ops[2])
+
+    # The global step should have been updated since we only need to collect 2
+    # gradients. The variables should now have the new values after the average
+    # of the gradients from worker 0/2 are applied.
+    while sessions[1].run(global_step) != 1:
+      time.sleep(0.01)
+
+    self.assertAllEqual(1, sessions[1].run(global_step))
+    self.assertAllClose(0 - (0.1 + 0.5) / 2 * 2.0, sessions[1].run(var_0_g_1))
+    self.assertAllClose(1 - (0.9 + 1.3) / 2 * 2.0, sessions[1].run(var_1_g_1))
+
+    # Worker 1 finished later and its gradients will now be dropped as it is
+    # stale.
+    sessions[1].run(train_ops[1])
+
+    # As shown in the previous test, the local_step for all workers should be
+    # still 0 so their next computation will also be dropped.
+    sessions[0].run(train_ops[0])
+    sessions[1].run(train_ops[1])
+    sessions[2].run(train_ops[2])
+
+    # Although the global step should still be 1 as explained above, the local
+    # step should now be updated to 1. Just check worker 1 as an example.
+    self.assertAllEqual(1, sessions[1].run(global_step))
+    self.assertAllEqual(1, sessions[1].run(local_step_1))
+
+    thread_0 = self.checkedThread(
+        target=self._run, args=(train_ops[0], sessions[0]))
+    thread_1 = self.checkedThread(
+        target=self._run, args=(train_ops[1], sessions[1]))
+
+    # Lets worker 0 execute first.
+    # It will wait as we need 2 workers to finish this step and the global step
+    # should be still 1.
+    thread_0.start()
+    self.assertAllEqual(1, sessions[1].run(global_step))
+
+    # Starts worker 1.
+    thread_1.start()
+    thread_1.join()
+    thread_0.join()
+
+    # The global step should now be 2 and the gradients should have been
+    # applied again.
+    self.assertAllEqual(2, sessions[1].run(global_step))
+    self.assertAllClose(-0.6 - (0.1 + 0.3) / 2 * 2.0,
+                        sessions[1].run(var_0_g_1))
+    self.assertAllClose(-1.2 - (0.9 + 1.1) / 2 * 2.0,
+                        sessions[1].run(var_1_g_1))
+
+
+class SyncReplicasOptimizerHookTest(test.TestCase):
+
+  def testErrorIfUsedBeforeMinimizeCalled(self):
+    opt = training.SyncReplicasOptimizer(
+        opt=gradient_descent.GradientDescentOptimizer(1.0),
+        replicas_to_aggregate=1,
+        total_num_replicas=1)
+    hook = opt.make_session_run_hook(True)
+    with self.assertRaisesRegex(ValueError, "apply_gradient should be called"):
+      hook.begin()
+
+  @test_util.run_v1_only(
+      "train.SyncReplicasOptimizer and train.GradientDescentOptimizer "
+      "are V1 only APIs.")
+  def testCanCreatedBeforeMinimizeCalled(self):
+    opt = training.SyncReplicasOptimizer(
+        opt=gradient_descent.GradientDescentOptimizer(1.0),
+        replicas_to_aggregate=1,
+        total_num_replicas=1)
+    hook = opt.make_session_run_hook(True)
+    v = variable_v1.VariableV1([0.])
+    global_step = variable_v1.VariableV1(0, name="global_step", trainable=False)
+    opt.minimize(v, global_step=global_step)
+    hook.begin()
+
+  @test_util.run_v1_only(
+      "train.SyncReplicasOptimizer and train.AdamOptimizer are V1 only APIs.")
+  def testFetchVariableList(self):
+    opt = training.SyncReplicasOptimizer(
+        opt=adam.AdamOptimizer(0.01),
+        replicas_to_aggregate=1,
+        total_num_replicas=1)
+    v = variable_v1.VariableV1([0.], name="fetch_variable_test")
+    global_step = variable_v1.VariableV1(0, name="global_step", trainable=False)
+    opt.minimize(v, global_step=global_step)
+    opt_variables = opt.variables()
+    beta1_power, beta2_power = opt._opt._get_beta_accumulators()
+    self.assertIn(beta1_power, opt_variables)
+    self.assertIn(beta2_power, opt_variables)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/training.py

@@ -0,0 +1,733 @@
+# 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.
+# ==============================================================================
+
+"""Support for training models.
+
+See the [Training](https://tensorflow.org/api_guides/python/train) guide.
+"""
+
+# Optimizers.
+# pylint: disable=g-bad-import-order,unused-import
+from tensorflow.python.ops.sdca_ops import sdca_optimizer
+from tensorflow.python.ops.sdca_ops import sdca_fprint
+from tensorflow.python.ops.sdca_ops import sdca_shrink_l1
+from tensorflow.python.training.adadelta import AdadeltaOptimizer
+from tensorflow.python.training.adagrad import AdagradOptimizer
+from tensorflow.python.training.adagrad_da import AdagradDAOptimizer
+from tensorflow.python.training.proximal_adagrad import ProximalAdagradOptimizer
+from tensorflow.python.training.adam import AdamOptimizer
+from tensorflow.python.training.ftrl import FtrlOptimizer
+from tensorflow.python.training.experimental.loss_scale_optimizer import MixedPrecisionLossScaleOptimizer
+from tensorflow.python.training.experimental.mixed_precision import enable_mixed_precision_graph_rewrite_v1
+from tensorflow.python.training.momentum import MomentumOptimizer
+from tensorflow.python.training.moving_averages import ExponentialMovingAverage
+from tensorflow.python.training.optimizer import Optimizer
+from tensorflow.python.training.rmsprop import RMSPropOptimizer
+from tensorflow.python.training.gradient_descent import GradientDescentOptimizer
+from tensorflow.python.training.proximal_gradient_descent import ProximalGradientDescentOptimizer
+from tensorflow.python.training.sync_replicas_optimizer import SyncReplicasOptimizer
+
+# Utility classes for training.
+from tensorflow.python.training.coordinator import Coordinator
+from tensorflow.python.training.coordinator import LooperThread
+# go/tf-wildcard-import
+# pylint: disable=wildcard-import
+from tensorflow.python.training.queue_runner import *
+
+# For the module level doc.
+from tensorflow.python.training import input as _input
+from tensorflow.python.training.input import *  # pylint: disable=redefined-builtin
+# pylint: enable=wildcard-import
+
+from tensorflow.python.training.basic_session_run_hooks import get_or_create_steps_per_run_variable
+from tensorflow.python.training.basic_session_run_hooks import SecondOrStepTimer
+from tensorflow.python.training.basic_session_run_hooks import LoggingTensorHook
+from tensorflow.python.training.basic_session_run_hooks import StopAtStepHook
+from tensorflow.python.training.basic_session_run_hooks import CheckpointSaverHook
+from tensorflow.python.training.basic_session_run_hooks import CheckpointSaverListener
+from tensorflow.python.training.basic_session_run_hooks import StepCounterHook
+from tensorflow.python.training.basic_session_run_hooks import NanLossDuringTrainingError
+from tensorflow.python.training.basic_session_run_hooks import NanTensorHook
+from tensorflow.python.training.basic_session_run_hooks import SummarySaverHook
+from tensorflow.python.training.basic_session_run_hooks import GlobalStepWaiterHook
+from tensorflow.python.training.basic_session_run_hooks import FinalOpsHook
+from tensorflow.python.training.basic_session_run_hooks import FeedFnHook
+from tensorflow.python.training.basic_session_run_hooks import ProfilerHook
+from tensorflow.python.training.basic_loops import basic_train_loop
+from tensorflow.python.trackable.python_state import PythonState
+from tensorflow.python.checkpoint.checkpoint import Checkpoint
+from tensorflow.python.checkpoint.checkpoint_view import CheckpointView
+from tensorflow.python.training.checkpoint_utils import init_from_checkpoint
+from tensorflow.python.training.checkpoint_utils import list_variables
+from tensorflow.python.training.checkpoint_utils import load_checkpoint
+from tensorflow.python.training.checkpoint_utils import load_variable
+
+from tensorflow.python.training.device_setter import replica_device_setter
+from tensorflow.python.training.monitored_session import Scaffold
+from tensorflow.python.training.monitored_session import MonitoredTrainingSession
+from tensorflow.python.training.monitored_session import SessionCreator
+from tensorflow.python.training.monitored_session import ChiefSessionCreator
+from tensorflow.python.training.monitored_session import WorkerSessionCreator
+from tensorflow.python.training.monitored_session import MonitoredSession
+from tensorflow.python.training.monitored_session import SingularMonitoredSession
+from tensorflow.python.training.saver import Saver
+from tensorflow.python.checkpoint.checkpoint_management import checkpoint_exists
+from tensorflow.python.checkpoint.checkpoint_management import generate_checkpoint_state_proto
+from tensorflow.python.checkpoint.checkpoint_management import get_checkpoint_mtimes
+from tensorflow.python.checkpoint.checkpoint_management import get_checkpoint_state
+from tensorflow.python.checkpoint.checkpoint_management import latest_checkpoint
+from tensorflow.python.checkpoint.checkpoint_management import update_checkpoint_state
+from tensorflow.python.training.saver import export_meta_graph
+from tensorflow.python.training.saver import import_meta_graph
+from tensorflow.python.training.saving import saveable_object_util
+from tensorflow.python.training.session_run_hook import SessionRunHook
+from tensorflow.python.training.session_run_hook import SessionRunArgs
+from tensorflow.python.training.session_run_hook import SessionRunContext
+from tensorflow.python.training.session_run_hook import SessionRunValues
+from tensorflow.python.training.session_manager import SessionManager
+from tensorflow.python.training.summary_io import summary_iterator
+from tensorflow.python.training.supervisor import Supervisor
+from tensorflow.python.training.training_util import write_graph
+from tensorflow.python.training.training_util import global_step
+from tensorflow.python.training.training_util import get_global_step
+from tensorflow.python.training.training_util import assert_global_step
+from tensorflow.python.training.training_util import create_global_step
+from tensorflow.python.training.training_util import get_or_create_global_step
+from tensorflow.python.training.warm_starting_util import VocabInfo
+from tensorflow.python.training.warm_starting_util import warm_start
+from tensorflow.python.training.py_checkpoint_reader import NewCheckpointReader
+from tensorflow.python.util.tf_export import tf_export
+
+# pylint: disable=wildcard-import
+# Training data protos.
+from tensorflow.core.example.example_pb2 import *
+from tensorflow.core.example.feature_pb2 import *
+from tensorflow.core.protobuf.saver_pb2 import *
+
+# Utility op.  Open Source. TODO(touts): move to nn?
+from tensorflow.python.training.learning_rate_decay import *
+# pylint: enable=wildcard-import
+
+# Distributed computing support.
+from tensorflow.core.protobuf.cluster_pb2 import ClusterDef
+from tensorflow.core.protobuf.cluster_pb2 import JobDef
+from tensorflow.core.protobuf.tensorflow_server_pb2 import ServerDef
+from tensorflow.python.training.server_lib import ClusterSpec
+from tensorflow.python.training.server_lib import Server
+
+# pylint: disable=undefined-variable
+tf_export("train.BytesList")(BytesList)
+tf_export("train.ClusterDef")(ClusterDef)
+tf_export("train.Example")(Example)
+tf_export("train.Feature")(Feature)
+tf_export("train.Features")(Features)
+tf_export("train.FeatureList")(FeatureList)
+tf_export("train.FeatureLists")(FeatureLists)
+tf_export("train.FloatList")(FloatList)
+tf_export("train.Int64List")(Int64List)
+tf_export("train.JobDef")(JobDef)
+tf_export(v1=["train.SaverDef"])(SaverDef)
+tf_export("train.SequenceExample")(SequenceExample)
+tf_export("train.ServerDef")(ServerDef)
+
+BytesList.__doc__ = """\
+Used in `tf.train.Example` protos. Holds a list of byte-strings.
+
+An `Example` proto is a representation of the following python type:
+
+```
+Dict[str,
+     Union[List[bytes],
+           List[int64],
+           List[float]]]
+```
+
+This proto implements the `List[bytes]` portion.
+
+>>> from google.protobuf import text_format
+>>> example = text_format.Parse('''
+...   features {
+...     feature {key: "my_feature"
+...              value {bytes_list {value: ['abc', '12345' ]}}}
+...   }''',
+...   tf.train.Example())
+>>>
+>>> example.features.feature['my_feature'].bytes_list.value
+["abc", "12345"]
+
+Use `tf.io.parse_example` to extract tensors from a serialized `Example` proto:
+
+>>> tf.io.parse_example(
+...     example.SerializeToString(),
+...     features = {'my_feature': tf.io.RaggedFeature(dtype=tf.string)})
+{'my_feature': <tf.Tensor: shape=(2,), dtype=string,
+                           numpy=array([b'abc', b'12345'], dtype=object)>}
+
+
+See the [`tf.train.Example`](https://www.tensorflow.org/tutorials/load_data/tfrecord#tftrainexample)
+guide for usage details.
+"""
+
+FloatList.__doc__ = """\
+Used in `tf.train.Example` protos. Holds a list of floats.
+
+An `Example` proto is a representation of the following python type:
+
+```
+Dict[str,
+     Union[List[bytes],
+           List[int64],
+           List[float]]]
+```
+
+This proto implements the `List[float]` portion.
+
+>>> from google.protobuf import text_format
+>>> example = text_format.Parse('''
+...   features {
+...     feature {key: "my_feature"
+...              value {float_list {value: [1., 2., 3., 4. ]}}}
+...   }''',
+...   tf.train.Example())
+>>>
+>>> example.features.feature['my_feature'].float_list.value
+[1.0, 2.0, 3.0, 4.0]
+
+Use `tf.io.parse_example` to extract tensors from a serialized `Example` proto:
+
+>>> tf.io.parse_example(
+...     example.SerializeToString(),
+...     features = {'my_feature': tf.io.RaggedFeature(dtype=tf.float32)})
+{'my_feature': <tf.Tensor: shape=(4,), dtype=float32,
+                           numpy=array([1., 2., 3., 4.], dtype=float32)>}
+
+See the [`tf.train.Example`](https://www.tensorflow.org/tutorials/load_data/tfrecord#tftrainexample)
+guide for usage details.
+"""
+
+Int64List.__doc__ = """\
+Used in `tf.train.Example` protos. Holds a list of Int64s.
+
+An `Example` proto is a representation of the following python type:
+
+```
+Dict[str,
+     Union[List[bytes],
+           List[int64],
+           List[float]]]
+```
+
+This proto implements the `List[int64]` portion.
+
+>>> from google.protobuf import text_format
+>>> example = text_format.Parse('''
+...   features {
+...     feature {key: "my_feature"
+...              value {int64_list {value: [1, 2, 3, 4]}}}
+...   }''',
+...   tf.train.Example())
+>>>
+>>> example.features.feature['my_feature'].int64_list.value
+[1, 2, 3, 4]
+
+Use `tf.io.parse_example` to extract tensors from a serialized `Example` proto:
+
+>>> tf.io.parse_example(
+...     example.SerializeToString(),
+...     features = {'my_feature': tf.io.RaggedFeature(dtype=tf.int64)})
+{'my_feature': <tf.Tensor: shape=(4,), dtype=float32,
+                           numpy=array([1, 2, 3, 4], dtype=int64)>}
+
+See the [`tf.train.Example`](https://www.tensorflow.org/tutorials/load_data/tfrecord#tftrainexample)
+guide for usage details.
+"""
+
+Feature.__doc__ = """\
+Used in `tf.train.Example` protos. Contains a list of values.
+
+An `Example` proto is a representation of the following python type:
+
+```
+Dict[str,
+     Union[List[bytes],
+           List[int64],
+           List[float]]]
+```
+
+This proto implements the `Union`.
+
+The contained list can be one of three types:
+
+  - `tf.train.BytesList`
+  - `tf.train.FloatList`
+  - `tf.train.Int64List`
+
+>>> int_feature = tf.train.Feature(
+...     int64_list=tf.train.Int64List(value=[1, 2, 3, 4]))
+>>> float_feature = tf.train.Feature(
+...     float_list=tf.train.FloatList(value=[1., 2., 3., 4.]))
+>>> bytes_feature = tf.train.Feature(
+...     bytes_list=tf.train.BytesList(value=[b"abc", b"1234"]))
+>>>
+>>> example = tf.train.Example(
+...     features=tf.train.Features(feature={
+...         'my_ints': int_feature,
+...         'my_floats': float_feature,
+...         'my_bytes': bytes_feature,
+...     }))
+
+Use `tf.io.parse_example` to extract tensors from a serialized `Example` proto:
+
+>>> tf.io.parse_example(
+...     example.SerializeToString(),
+...     features = {
+...         'my_ints': tf.io.RaggedFeature(dtype=tf.int64),
+...         'my_floats': tf.io.RaggedFeature(dtype=tf.float32),
+...         'my_bytes': tf.io.RaggedFeature(dtype=tf.string)})
+{'my_bytes': <tf.Tensor: shape=(2,), dtype=string,
+                         numpy=array([b'abc', b'1234'], dtype=object)>,
+ 'my_floats': <tf.Tensor: shape=(4,), dtype=float32,
+                          numpy=array([1., 2., 3., 4.], dtype=float32)>,
+ 'my_ints': <tf.Tensor: shape=(4,), dtype=int64,
+                        numpy=array([1, 2, 3, 4])>}
+
+"""
+
+Features.__doc__ = """\
+Used in `tf.train.Example` protos. Contains the mapping from keys to `Feature`.
+
+An `Example` proto is a representation of the following python type:
+
+```
+Dict[str,
+     Union[List[bytes],
+           List[int64],
+           List[float]]]
+```
+
+This proto implements the `Dict`.
+
+>>> int_feature = tf.train.Feature(
+...     int64_list=tf.train.Int64List(value=[1, 2, 3, 4]))
+>>> float_feature = tf.train.Feature(
+...     float_list=tf.train.FloatList(value=[1., 2., 3., 4.]))
+>>> bytes_feature = tf.train.Feature(
+...     bytes_list=tf.train.BytesList(value=[b"abc", b"1234"]))
+>>>
+>>> example = tf.train.Example(
+...     features=tf.train.Features(feature={
+...         'my_ints': int_feature,
+...         'my_floats': float_feature,
+...         'my_bytes': bytes_feature,
+...     }))
+
+Use `tf.io.parse_example` to extract tensors from a serialized `Example` proto:
+
+>>> tf.io.parse_example(
+...     example.SerializeToString(),
+...     features = {
+...         'my_ints': tf.io.RaggedFeature(dtype=tf.int64),
+...         'my_floats': tf.io.RaggedFeature(dtype=tf.float32),
+...         'my_bytes': tf.io.RaggedFeature(dtype=tf.string)})
+{'my_bytes': <tf.Tensor: shape=(2,), dtype=string,
+                         numpy=array([b'abc', b'1234'], dtype=object)>,
+ 'my_floats': <tf.Tensor: shape=(4,), dtype=float32,
+                          numpy=array([1., 2., 3., 4.], dtype=float32)>,
+ 'my_ints': <tf.Tensor: shape=(4,), dtype=int64,
+                        numpy=array([1, 2, 3, 4])>}
+
+"""
+
+FeatureList.__doc__ = """\
+Mainly used as part of a `tf.train.SequenceExample`.
+
+Contains a list of `tf.train.Feature`s.
+
+The `tf.train.SequenceExample` proto can be thought of as a
+proto implementation of the following python type:
+
+```
+# tf.train.Feature
+Feature = Union[List[bytes],
+                List[int64],
+                List[float]]
+
+# tf.train.FeatureList
+FeatureList = List[Feature]
+
+# tf.train.FeatureLists
+FeatureLists = Dict[str, FeatureList]
+
+class SequenceExample(typing.NamedTuple):
+  context: Dict[str, Feature]
+  feature_lists: FeatureLists
+```
+
+This proto implements the `List[Feature]` portion.
+
+"""
+
+FeatureLists.__doc__ = """\
+Mainly used as part of a `tf.train.SequenceExample`.
+
+Contains a list of `tf.train.Feature`s.
+
+The `tf.train.SequenceExample` proto can be thought of as a
+proto implementation of the following python type:
+
+```
+# tf.train.Feature
+Feature = Union[List[bytes],
+                List[int64],
+                List[float]]
+
+# tf.train.FeatureList
+FeatureList = List[Feature]
+
+# tf.train.FeatureLists
+FeatureLists = Dict[str, FeatureList]
+
+class SequenceExample(typing.NamedTuple):
+  context: Dict[str, Feature]
+  feature_lists: FeatureLists
+```
+
+This proto implements the `Dict[str, FeatureList]` portion.
+"""
+
+
+Example.__doc__ = """\
+An `Example` is a standard proto storing data for training and inference.
+
+An `Example` proto is a representation of the following python type:
+
+```
+Dict[str,
+     Union[List[bytes],
+           List[int64],
+           List[float]]]
+```
+
+It contains a key-value store `Example.features` where each key (string) maps
+to a `tf.train.Feature` message which contains a fixed-type list. This flexible
+and compact format allows the storage of large amounts of typed data, but
+requires that the data shape and use be determined by the configuration files
+and parsers that are used to read and write this format (refer to
+`tf.io.parse_example` for details).
+
+>>> from google.protobuf import text_format
+>>> example = text_format.Parse('''
+...   features {
+...     feature {key: "my_feature"
+...              value {int64_list {value: [1, 2, 3, 4]}}}
+...   }''',
+...   tf.train.Example())
+
+Use `tf.io.parse_example` to extract tensors from a serialized `Example` proto:
+
+>>> tf.io.parse_example(
+...     example.SerializeToString(),
+...     features = {'my_feature': tf.io.RaggedFeature(dtype=tf.int64)})
+{'my_feature': <tf.Tensor: shape=(4,), dtype=float32,
+                           numpy=array([1, 2, 3, 4], dtype=int64)>}
+
+While the list of keys, and the contents of each key _could_ be different for
+every `Example`, TensorFlow expects a fixed list of keys, each with a fixed
+`tf.dtype`. A conformant `Example` dataset obeys the following conventions:
+
+  - If a Feature `K` exists in one example with data type `T`, it must be of
+      type `T` in all other examples when present. It may be omitted.
+  - The number of instances of Feature `K` list data may vary across examples,
+      depending on the requirements of the model.
+  - If a Feature `K` doesn't exist in an example, a `K`-specific default will be
+      used, if configured.
+  - If a Feature `K` exists in an example but contains no items, the intent
+      is considered to be an empty tensor and no default will be used.
+
+"""
+
+SequenceExample.__doc__ = """\
+A `SequenceExample` represents a sequence of features and some context.
+
+It can be thought of as a proto-implementation of the following python type:
+
+```
+Feature = Union[List[bytes],
+                List[int64],
+                List[float]]
+
+class SequenceExample(typing.NamedTuple):
+  context: Dict[str, Feature]
+  feature_lists: Dict[str, List[Feature]]
+```
+
+To implement this as protos it's broken up into sub-messages as follows:
+
+```
+# tf.train.Feature
+Feature = Union[List[bytes],
+                List[int64],
+                List[float]]
+
+# tf.train.FeatureList
+FeatureList = List[Feature]
+
+# tf.train.FeatureLists
+FeatureLists = Dict[str, FeatureList]
+
+# tf.train.SequenceExample
+class SequenceExample(typing.NamedTuple):
+  context: Dict[str, Feature]
+  feature_lists: FeatureLists
+```
+
+To parse a `SequenceExample` in TensorFlow refer to the
+`tf.io.parse_sequence_example` function.
+
+The `context` contains features which apply to the entire
+example. The `feature_lists` contain a key, value map where each key is
+associated with a repeated set of `tf.train.Features` (a `tf.train.FeatureList`).
+A `FeatureList` represents the values of a feature identified by its key
+over time / frames.
+
+Below is a `SequenceExample` for a movie recommendation application recording a
+sequence of ratings by a user. The time-independent features ("locale",
+"age", "favorites") describing the user are part of the context. The sequence
+of movies the user rated are part of the feature_lists. For each movie in the
+sequence we have information on its name and actors and the user's rating.
+This information is recorded in three separate `feature_list`s.
+In the example below there are only two movies. All three `feature_list`s,
+namely "movie_ratings", "movie_names", and "actors" have a feature value for
+both movies. Note, that "actors" is itself a `bytes_list` with multiple
+strings per movie.
+
+```
+  context: {
+    feature: {
+      key  : "locale"
+      value: {
+        bytes_list: {
+          value: [ "pt_BR" ]
+        }
+      }
+    }
+    feature: {
+      key  : "age"
+      value: {
+        float_list: {
+          value: [ 19.0 ]
+        }
+      }
+    }
+    feature: {
+      key  : "favorites"
+      value: {
+        bytes_list: {
+          value: [ "Majesty Rose", "Savannah Outen", "One Direction" ]
+        }
+      }
+    }
+  }
+  feature_lists: {
+    feature_list: {
+      key  : "movie_ratings"
+      value: {
+        feature: {
+          float_list: {
+            value: [ 4.5 ]
+          }
+        }
+        feature: {
+          float_list: {
+            value: [ 5.0 ]
+          }
+        }
+      }
+    }
+    feature_list: {
+      key  : "movie_names"
+      value: {
+        feature: {
+          bytes_list: {
+            value: [ "The Shawshank Redemption" ]
+          }
+        }
+        feature: {
+          bytes_list: {
+            value: [ "Fight Club" ]
+          }
+        }
+      }
+    }
+    feature_list: {
+      key  : "actors"
+      value: {
+        feature: {
+          bytes_list: {
+            value: [ "Tim Robbins", "Morgan Freeman" ]
+          }
+        }
+        feature: {
+          bytes_list: {
+            value: [ "Brad Pitt", "Edward Norton", "Helena Bonham Carter" ]
+          }
+        }
+      }
+    }
+  }
+```
+
+A conformant `SequenceExample` data set obeys the following conventions:
+
+`context`:
+
+  - All conformant context features `K` must obey the same conventions as
+    a conformant Example's features (see above).
+
+`feature_lists`:
+
+  - A `FeatureList L` may be missing in an example; it is up to the
+    parser configuration to determine if this is allowed or considered
+    an empty list (zero length).
+  - If a `FeatureList L` exists, it may be empty (zero length).
+  - If a `FeatureList L` is non-empty, all features within the `FeatureList`
+    must have the same data type `T`. Even across `SequenceExample`s, the type `T`
+    of the `FeatureList` identified by the same key must be the same. An entry
+    without any values may serve as an empty feature.
+  - If a `FeatureList L` is non-empty, it is up to the parser configuration
+    to determine if all features within the `FeatureList` must
+    have the same size.  The same holds for this `FeatureList` across multiple
+    examples.
+  - For sequence modeling ([example](https://github.com/tensorflow/nmt)), the
+    feature lists represent a sequence of frames. In this scenario, all
+    `FeatureList`s in a `SequenceExample` have the same number of `Feature`
+    messages, so that the i-th element in each `FeatureList` is part of the
+    i-th frame (or time step).
+
+**Examples of conformant and non-conformant examples' `FeatureLists`:**
+
+Conformant `FeatureLists`:
+
+```
+    feature_lists: { feature_list: {
+      key: "movie_ratings"
+      value: { feature: { float_list: { value: [ 4.5 ] } }
+               feature: { float_list: { value: [ 5.0 ] } } }
+    } }
+```
+
+Non-conformant `FeatureLists` (mismatched types):
+
+```
+    feature_lists: { feature_list: {
+      key: "movie_ratings"
+      value: { feature: { float_list: { value: [ 4.5 ] } }
+               feature: { int64_list: { value: [ 5 ] } } }
+    } }
+```
+
+Conditionally conformant `FeatureLists`, the parser configuration determines
+if the feature sizes must match:
+
+```
+    feature_lists: { feature_list: {
+      key: "movie_ratings"
+      value: { feature: { float_list: { value: [ 4.5 ] } }
+               feature: { float_list: { value: [ 5.0, 6.0 ] } } }
+    } }
+```
+
+**Examples of conformant and non-conformant `SequenceExample`s:**
+
+Conformant pair of SequenceExample:
+
+```
+    feature_lists: { feature_list: {
+      key: "movie_ratings"
+      value: { feature: { float_list: { value: [ 4.5 ] } }
+               feature: { float_list: { value: [ 5.0 ] } } }
+     } }
+
+    feature_lists: { feature_list: {
+      key: "movie_ratings"
+      value: { feature: { float_list: { value: [ 4.5 ] } }
+               feature: { float_list: { value: [ 5.0 ] } }
+               feature: { float_list: { value: [ 2.0 ] } } }
+     } }
+```
+
+Conformant pair of `SequenceExample`s:
+
+```
+    feature_lists: { feature_list: {
+      key: "movie_ratings"
+      value: { feature: { float_list: { value: [ 4.5 ] } }
+               feature: { float_list: { value: [ 5.0 ] } } }
+     } }
+
+    feature_lists: { feature_list: {
+      key: "movie_ratings"
+      value: { }
+     } }
+```
+
+Conditionally conformant pair of `SequenceExample`s, the parser configuration
+determines if the second `feature_lists` is consistent (zero-length) or
+invalid (missing "movie_ratings"):
+
+```
+    feature_lists: { feature_list: {
+      key: "movie_ratings"
+      value: { feature: { float_list: { value: [ 4.5 ] } }
+               feature: { float_list: { value: [ 5.0 ] } } }
+     } }
+
+   feature_lists: { }
+```
+
+Non-conformant pair of `SequenceExample`s (mismatched types):
+
+```
+    feature_lists: { feature_list: {
+      key: "movie_ratings"
+      value: { feature: { float_list: { value: [ 4.5 ] } }
+               feature: { float_list: { value: [ 5.0 ] } } }
+     } }
+
+    feature_lists: { feature_list: {
+      key: "movie_ratings"
+      value: { feature: { int64_list: { value: [ 4 ] } }
+               feature: { int64_list: { value: [ 5 ] } }
+               feature: { int64_list: { value: [ 2 ] } } }
+     } }
+```
+
+Conditionally conformant pair of `SequenceExample`s; the parser configuration
+determines if the feature sizes must match:
+
+```
+    feature_lists: { feature_list: {
+      key: "movie_ratings"
+      value: { feature: { float_list: { value: [ 4.5 ] } }
+               feature: { float_list: { value: [ 5.0 ] } } }
+    } }
+
+    feature_lists: { feature_list: {
+      key: "movie_ratings"
+      value: { feature: { float_list: { value: [ 4.0 ] } }
+              feature: { float_list: { value: [ 5.0, 3.0 ] } }
+    } }
+```
+"""

tensorflow/python/training/training_ops_test.py

@@ -0,0 +1,511 @@
+# 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.
+# ==============================================================================
+"""Tests for tensorflow.ops.gen_training_ops."""
+
+import itertools
+import threading
+
+import numpy as np
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.framework.test_util import TensorFlowTestCase
+# Import resource_variable_ops for the variables-to-tensor implicit conversion.
+from tensorflow.python.ops import gen_training_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops  # pylint: disable=unused-import
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import googletest
+
+
+class TrainingOpsTest(TensorFlowTestCase):
+
+  def _toType(self, dtype):
+    if dtype == np.float16:
+      return dtypes.float16
+    elif dtype == np.float32:
+      return dtypes.float32
+    elif dtype == np.float64:
+      return dtypes.float64
+    elif dtype == np.int32:
+      return dtypes.int32
+    elif dtype == np.int64:
+      return dtypes.int64
+    else:
+      assert False, (dtype)
+
+  def _testTypes(self, x, alpha, delta, use_gpu=None):
+    self.setUp()
+    with self.session(use_gpu=use_gpu):
+      var = variable_v1.VariableV1(x)
+      self.evaluate(variables.global_variables_initializer())
+      self.assertAllCloseAccordingToType(x, self.evaluate(var))
+      apply_sgd = gen_training_ops.apply_gradient_descent(var, alpha, delta)
+      out = self.evaluate(apply_sgd)
+      self.assertShapeEqual(out, apply_sgd)
+      self.assertAllCloseAccordingToType(x - alpha * delta, out)
+
+  @test_util.run_v1_only("ApplyGradientDescent op returns a ref, so it is not "
+                         "supported in eager mode.")
+  def testApplyGradientDescent(self):
+    for (dtype, use_gpu) in itertools.product(
+        [np.float16, np.float32, np.float64], [False, True]):
+      x = np.arange(100).astype(dtype)
+      alpha = np.array(2.0).astype(dtype)
+      delta = np.arange(100).astype(dtype)
+      self._testTypes(x, alpha, delta, use_gpu)
+
+  def _testTypesForAdagrad(self, x, y, lr, grad, use_gpu=None):
+    self.setUp()
+    with self.session(use_gpu=use_gpu):
+      var = variable_v1.VariableV1(x)
+      accum = variable_v1.VariableV1(y)
+      self.evaluate(variables.global_variables_initializer())
+
+      self.assertAllCloseAccordingToType(x, self.evaluate(var))
+      apply_adagrad = gen_training_ops.apply_adagrad(var, accum, lr, grad)
+      out = self.evaluate(apply_adagrad)
+      self.assertShapeEqual(out, apply_adagrad)
+      self.assertAllCloseAccordingToType(x - lr * grad * (y + grad * grad)**
+                                         (-0.5), out)
+      self.assertAllCloseAccordingToType(y + grad * grad, self.evaluate(accum))
+
+  def _testTypesForFtrl(self,
+                        x,
+                        y,
+                        z,
+                        lr,
+                        grad,
+                        use_gpu=None,
+                        l1=0.0,
+                        l2=0.0,
+                        lr_power=-0.5):
+    self.setUp()
+    with self.session(use_gpu=use_gpu):
+      var = variable_v1.VariableV1(x)
+      accum = variable_v1.VariableV1(y)
+      linear = variable_v1.VariableV1(z)
+      self.evaluate(variables.global_variables_initializer())
+
+      self.assertAllCloseAccordingToType(x, self.evaluate(var))
+      apply_ftrl = gen_training_ops.apply_ftrl(var, accum, linear, grad, lr, l1,
+                                               l2, lr_power)
+      out = self.evaluate(apply_ftrl)
+      self.assertShapeEqual(out, apply_ftrl)
+      accum_update = y + grad * grad
+      linear_update = z + grad - (accum_update**(-lr_power) - y**
+                                  (-lr_power)) / lr * x
+      quadratic = 1.0 / (accum_update**(lr_power) * lr) + 2 * l2
+      expected_out = np.array([(
+          np.sign(linear_update[i]) * l1 - linear_update[i]) / (quadratic[i]) if
+                               np.abs(linear_update[i]) > l1 else 0.0
+                               for i in range(linear_update.size)])
+      self.assertAllCloseAccordingToType(accum_update, self.evaluate(accum))
+      if x.dtype == np.float16:
+        # The calculations here really are not very precise in float16.
+        self.assertAllClose(
+            linear_update, self.evaluate(linear), rtol=2e-2, atol=2e-2)
+        self.assertAllClose(expected_out, out, rtol=2e-2, atol=2e-2)
+      elif x.dtype == np.float32:
+        # The calculations here not sufficiently precise in float32.
+        self.assertAllClose(
+            linear_update, self.evaluate(linear), rtol=1e-5, atol=1e-5)
+        self.assertAllClose(expected_out, out, rtol=1e-5, atol=1e-5)
+      else:
+        self.assertAllClose(linear_update, self.evaluate(linear))
+        self.assertAllClose(expected_out, out)
+
+  def _testTypesForFtrlMultiplyLinearByLr(self,
+                                          x,
+                                          y,
+                                          z,
+                                          lr,
+                                          grad,
+                                          use_gpu=None,
+                                          l1=0.0,
+                                          l2=0.0,
+                                          lr_power=-0.5):
+    self.setUp()
+    with self.session(use_gpu=use_gpu):
+      var = variable_v1.VariableV1(x)
+      accum = variable_v1.VariableV1(y)
+      linear = variable_v1.VariableV1(z)
+      self.evaluate(variables.global_variables_initializer())
+
+      self.assertAllCloseAccordingToType(x, self.evaluate(var))
+      apply_ftrl = (
+          gen_training_ops.apply_ftrl(
+              var,
+              accum,
+              linear,
+              grad,
+              lr,
+              l1,
+              l2,
+              lr_power,
+              multiply_linear_by_lr=True))
+      out = self.evaluate(apply_ftrl)
+      self.assertShapeEqual(out, apply_ftrl)
+      accum_update = y + grad * grad
+      linear_update = z + grad * lr - (accum_update**(-lr_power) - y**
+                                       (-lr_power)) * x
+      quadratic = accum_update**(-lr_power) + 2 * l2 * lr
+      expected_out = np.array([
+          (np.sign(linear_update[i]) * l1 * lr - linear_update[i]) /
+          (quadratic[i]) if np.abs(linear_update[i]) > l1 * lr else 0.0
+          for i in range(linear_update.size)
+      ])
+      self.assertAllCloseAccordingToType(accum_update, self.evaluate(accum))
+      if x.dtype == np.float16:
+        # The calculations here really are not very precise in float16.
+        self.assertAllClose(
+            linear_update, self.evaluate(linear), rtol=2e-2, atol=2e-2)
+        self.assertAllClose(expected_out, out, rtol=2e-2, atol=2e-2)
+      elif x.dtype == np.float32:
+        # The calculations here not sufficiently precise in float32.
+        self.assertAllClose(
+            linear_update, self.evaluate(linear), rtol=1e-5, atol=1e-5)
+        self.assertAllClose(expected_out, out, rtol=1e-5, atol=1e-5)
+      else:
+        self.assertAllClose(linear_update, self.evaluate(linear))
+        self.assertAllClose(expected_out, out)
+
+  @test_util.run_v1_only("ApplyAdagrad op returns a ref, so it is not "
+                         "supported in eager mode.")
+  def testApplyAdagrad(self):
+    for (dtype, use_gpu) in itertools.product(
+        [np.float16, np.float32, np.float64], [False, True]):
+      x = np.arange(100).astype(dtype)
+      y = np.arange(1, 101).astype(dtype)
+      lr = np.array(2.0).astype(dtype)
+      grad = np.arange(100).astype(dtype)
+      self._testTypesForAdagrad(x, y, lr, grad, use_gpu)
+
+  @test_util.run_v1_only("ApplyFtrl op returns a ref, so it is not "
+                         "supported in eager mode.")
+  def testApplyFtrl(self):
+    for dtype in [np.float16, np.float32, np.float64]:
+      x = np.arange(100).astype(dtype)
+      y = np.arange(1, 101).astype(dtype)
+      z = np.arange(102, 202).astype(dtype)
+      lr = np.array(2.0).astype(dtype)
+      l1 = np.array(3.0).astype(dtype)
+      l2 = np.array(4.0).astype(dtype)
+      grad = np.arange(100).astype(dtype)
+      self._testTypesForFtrl(x, y, z, lr, grad, use_gpu=False, l1=l1, l2=l2)
+
+  @test_util.run_v1_only("ApplyFtrlMultiplyLinearByLr op returns a ref, so it "
+                         "is not supported in eager mode.")
+  def testApplyFtrlMultiplyLinearByLr(self):
+    for dtype in [np.float16, np.float32, np.float64]:
+      x = np.arange(100).astype(dtype)
+      y = np.arange(1, 101).astype(dtype)
+      z = np.arange(102, 202).astype(dtype)
+      lr = np.array(2.0).astype(dtype)
+      l1 = np.array(3.0).astype(dtype)
+      l2 = np.array(4.0).astype(dtype)
+      grad = np.arange(100).astype(dtype)
+      self._testTypesForFtrlMultiplyLinearByLr(
+          x, y, z, lr, grad, use_gpu=False, l1=l1, l2=l2)
+
+  def _testTypesForSparseAdagrad(self, x, y, lr, grad, indices, use_gpu):
+    self.setUp()
+    with self.session(use_gpu=use_gpu):
+      var = variable_v1.VariableV1(x)
+      accum = variable_v1.VariableV1(y)
+      self.evaluate(variables.global_variables_initializer())
+
+      self.assertAllCloseAccordingToType(x, self.evaluate(var))
+      sparse_apply_adagrad = gen_training_ops.sparse_apply_adagrad(
+          var, accum, lr, grad,
+          constant_op.constant(indices, self._toType(indices.dtype)))
+      out = self.evaluate(sparse_apply_adagrad)
+      self.assertShapeEqual(out, sparse_apply_adagrad)
+
+      for (i, index) in enumerate(indices):
+        self.assertAllCloseAccordingToType(
+            x[index] - lr * grad[i] * (y[index] + grad[i] * grad[i])**(-0.5),
+            self.evaluate(var)[index])
+        self.assertAllCloseAccordingToType(y[index] + grad[i] * grad[i],
+                                           self.evaluate(accum)[index])
+
+  def _testTypesForSparseFtrl(self,
+                              x,
+                              y,
+                              z,
+                              lr,
+                              grad,
+                              indices,
+                              use_gpu,
+                              l1=0.0,
+                              l2=0.0,
+                              lr_power=-0.5):
+    self.setUp()
+    with self.session(use_gpu=use_gpu):
+      var = variable_v1.VariableV1(x)
+      accum = variable_v1.VariableV1(y)
+      linear = variable_v1.VariableV1(z)
+      self.evaluate(variables.global_variables_initializer())
+
+      self.assertAllCloseAccordingToType(x, self.evaluate(var))
+      sparse_apply_ftrl = gen_training_ops.sparse_apply_ftrl(
+          var,
+          accum,
+          linear,
+          grad,
+          constant_op.constant(indices, self._toType(indices.dtype)),
+          lr,
+          l1,
+          l2,
+          lr_power=lr_power)
+      out = self.evaluate(sparse_apply_ftrl)
+      self.assertShapeEqual(out, sparse_apply_ftrl)
+
+      for (i, index) in enumerate(indices):
+        self.assertAllCloseAccordingToType(
+            x[index] - lr * grad[i] *
+            (y[index] + grad[i] * grad[i])**(lr_power),
+            self.evaluate(var)[index])
+        self.assertAllCloseAccordingToType(y[index] + grad[i] * grad[i],
+                                           self.evaluate(accum)[index])
+
+  def _testTypesForSparseFtrlMultiplyLinearByLr(self,
+                                                x,
+                                                y,
+                                                z,
+                                                lr,
+                                                grad,
+                                                indices,
+                                                l1=0.0,
+                                                l2=0.0,
+                                                lr_power=-0.5):
+    self.setUp()
+    with self.session(use_gpu=False):
+      var = variable_v1.VariableV1(x)
+      accum = variable_v1.VariableV1(y)
+      linear = variable_v1.VariableV1(z)
+      self.evaluate(variables.global_variables_initializer())
+
+      self.assertAllCloseAccordingToType(x, self.evaluate(var))
+      sparse_apply_ftrl = (
+          gen_training_ops.sparse_apply_ftrl(
+              var,
+              accum,
+              linear,
+              grad,
+              constant_op.constant(indices, self._toType(indices.dtype)),
+              lr,
+              l1,
+              l2,
+              lr_power=lr_power,
+              multiply_linear_by_lr=True))
+      out = self.evaluate(sparse_apply_ftrl)
+      self.assertShapeEqual(out, sparse_apply_ftrl)
+
+      for (i, index) in enumerate(indices):
+        self.assertAllCloseAccordingToType(
+            x[index] - lr * grad[i] * (y[index] + grad[i] * grad[i])**
+            (lr_power),
+            self.evaluate(var)[index])
+        self.assertAllCloseAccordingToType(y[index] + grad[i] * grad[i],
+                                           self.evaluate(accum)[index])
+
+  @test_util.run_v1_only("SparseApplyAdagrad op returns a ref, so it is not "
+                         "supported in eager mode.")
+  def testSparseApplyAdagrad(self):
+    for (dtype, index_type,
+         use_gpu) in itertools.product([np.float16, np.float32, np.float64],
+                                       [np.int32, np.int64], [False, True]):
+      x_val = [np.arange(10), np.arange(10, 20), np.arange(20, 30)]
+      y_val = [np.arange(1, 11), np.arange(11, 21), np.arange(21, 31)]
+      x = np.array(x_val).astype(dtype)
+      y = np.array(y_val).astype(dtype)
+      lr = np.array(2.0).astype(dtype)
+      grad_val = [np.arange(10), np.arange(10)]
+      grad = np.array(grad_val).astype(dtype)
+      indices = np.array([0, 2]).astype(index_type)
+      self._testTypesForSparseAdagrad(x, y, lr, grad, indices, use_gpu)
+      # Empty sparse gradients.
+      empty_grad = np.zeros([0, 10], dtype=dtype)
+      empty_indices = np.zeros([0], dtype=index_type)
+      self._testTypesForSparseAdagrad(x, y, lr, empty_grad, empty_indices,
+                                      use_gpu)
+
+  @test_util.run_v1_only("SparseApplyAdagrad op returns a ref, so it is not "
+                         "supported in eager mode.")
+  def testSparseApplyAdagradDim1(self):
+    for (dtype, index_type,
+         use_gpu) in itertools.product([np.float16, np.float32, np.float64],
+                                       [np.int32, np.int64], [False, True]):
+      x_val = [[1.0], [2.0], [3.0]]
+      y_val = [[4.0], [5.0], [6.0]]
+      x = np.array(x_val).astype(dtype)
+      y = np.array(y_val).astype(dtype)
+      lr = np.array(2.0).astype(dtype)
+      grad_val = [[1.5], [2.5]]
+      grad = np.array(grad_val).astype(dtype)
+      indices = np.array([0, 2]).astype(index_type)
+      self._testTypesForSparseAdagrad(x, y, lr, grad, indices, use_gpu)
+
+  @test_util.run_v1_only("SparseApplyFtrl op returns a ref, so it is not "
+                         "supported in eager mode.")
+  def testSparseApplyFtrlDim1(self):
+    for (dtype, index_type,
+         use_gpu) in itertools.product([np.float16, np.float32, np.float64],
+                                       [np.int32, np.int64], [False, True]):
+      x_val = [[0.0], [0.0], [0.0]]
+      y_val = [[4.0], [5.0], [6.0]]
+      z_val = [[0.0], [0.0], [0.0]]
+      x = np.array(x_val).astype(dtype)
+      y = np.array(y_val).astype(dtype)
+      z = np.array(z_val).astype(dtype)
+      lr = np.array(2.0).astype(dtype)
+      grad_val = [[1.5], [2.5]]
+      grad = np.array(grad_val).astype(dtype)
+      indices = np.array([0, 2]).astype(index_type)
+      self._testTypesForSparseFtrl(x, y, z, lr, grad, indices, use_gpu)
+      # Empty sparse gradients.
+      empty_grad = np.zeros([0, 1], dtype=dtype)
+      empty_indices = np.zeros([0], dtype=index_type)
+      self._testTypesForSparseFtrl(x, y, z, lr, empty_grad, empty_indices,
+                                   use_gpu)
+
+  @test_util.run_v1_only("SparseApplyFtrlMultiplyLinearByLr op returns a ref, "
+                         "so it is not supported in eager mode.")
+  def testSparseApplyFtrlMultiplyLinearByLrDim1(self):
+    for (dtype,
+         index_type) in itertools.product([np.float16, np.float32, np.float64],
+                                          [np.int32, np.int64]):
+      x_val = [[0.0], [0.0], [0.0]]
+      y_val = [[4.0], [5.0], [6.0]]
+      z_val = [[0.0], [0.0], [0.0]]
+      x = np.array(x_val).astype(dtype)
+      y = np.array(y_val).astype(dtype)
+      z = np.array(z_val).astype(dtype)
+      lr = np.array(2.0).astype(dtype)
+      grad_val = [[1.5], [2.5]]
+      grad = np.array(grad_val).astype(dtype)
+      indices = np.array([0, 2]).astype(index_type)
+      self._testTypesForSparseFtrlMultiplyLinearByLr(x, y, z, lr, grad, indices)
+
+  @test_util.run_v1_only("ApplyAdam op returns a ref, so it is not "
+                         "supported in eager mode.")
+  def testApplyAdam(self):
+    for dtype, use_gpu in itertools.product(
+        [np.float16, np.float32, np.float64], [False, True]):
+      var = np.arange(100).astype(dtype)
+      m = np.arange(1, 101).astype(dtype)
+      v = np.arange(101, 201).astype(dtype)
+      grad = np.arange(100).astype(dtype)
+      self._testTypesForAdam(var, m, v, grad, use_gpu)
+
+  def _testTypesForAdam(self, var, m, v, grad, use_gpu):
+    self.setUp()
+    with self.session(use_gpu=use_gpu):
+      var_t = variable_v1.VariableV1(var)
+      m_t = variable_v1.VariableV1(m)
+      v_t = variable_v1.VariableV1(v)
+
+      t = 1
+      beta1 = np.array(0.9, dtype=var.dtype)
+      beta2 = np.array(0.999, dtype=var.dtype)
+      beta1_power = beta1**t
+      beta2_power = beta2**t
+      lr = np.array(0.001, dtype=var.dtype)
+      epsilon = np.array(1e-8, dtype=var.dtype)
+      beta1_t = constant_op.constant(beta1, self._toType(var.dtype), [])
+      beta2_t = constant_op.constant(beta2, self._toType(var.dtype), [])
+      beta1_power_t = variable_v1.VariableV1(beta1_power)
+      beta2_power_t = variable_v1.VariableV1(beta2_power)
+      lr_t = constant_op.constant(lr, self._toType(var.dtype), [])
+      epsilon_t = constant_op.constant(epsilon, self._toType(var.dtype), [])
+      self.evaluate(variables.global_variables_initializer())
+
+      self.assertAllCloseAccordingToType(var, self.evaluate(var_t))
+      new_var, _, _ = self._adamUpdateNumpy(var, grad, t, m, v, lr, beta1,
+                                            beta2, epsilon)
+      apply_adam = gen_training_ops.apply_adam(var_t, m_t, v_t, beta1_power_t,
+                                               beta2_power_t, lr_t, beta1_t,
+                                               beta2_t, epsilon_t, grad)
+      out = self.evaluate(apply_adam)
+      self.assertShapeEqual(out, apply_adam)
+      self.assertAllCloseAccordingToType(new_var, out)
+
+  def _adamUpdateNumpy(self, param, g_t, t, m, v, alpha, beta1, beta2, epsilon):
+    alpha_t = alpha * np.sqrt(1 - beta2**t) / (1 - beta1**t)
+
+    m_t = beta1 * m + (1 - beta1) * g_t
+    v_t = beta2 * v + (1 - beta2) * g_t * g_t
+
+    param_t = param - alpha_t * m_t / (np.sqrt(v_t) + epsilon)
+    return param_t, m_t, v_t
+
+  @test_util.run_v2_only
+  def testResourceSparseApplyAdagradV2AndDisableCopyOnReadRace(self):
+    dtype = np.float32
+    index_type = np.int32
+    x_val = [np.arange(10), np.arange(10, 20), np.arange(20, 30)]
+    y_val = [np.arange(1, 11), np.arange(11, 21), np.arange(21, 31)]
+    x = np.array(x_val).astype(dtype)
+    y = np.array(y_val).astype(dtype)
+    lr = np.array(0.001, dtype=dtype)
+    epsilon = np.array(1e-8, dtype=dtype)
+    grad_val = [np.arange(10), np.arange(10)]
+    grad = np.array(grad_val).astype(dtype)
+    indices = np.array([0, 2]).astype(index_type)
+    var = variables.Variable(x)
+    accum = variables.Variable(y)
+    num_iter = 1000
+    self.evaluate(variables.global_variables_initializer())
+
+    @def_function.function
+    def fn_disable_copy_on_read():
+      ret = constant_op.constant(0, dtypes.int32)
+      for i in math_ops.range(num_iter):
+        op1 = resource_variable_ops.disable_copy_on_read(var.handle)
+        op2 = resource_variable_ops.disable_copy_on_read(accum.handle)
+        with ops.control_dependencies([op1, op2]):
+          ret += i
+      return ret
+
+    @def_function.function
+    def fn_resource_sparse_apply_adagrad_v2():
+      ret = constant_op.constant(0, dtypes.int32)
+      for i in math_ops.range(num_iter):
+        adagrad_op = gen_training_ops.resource_sparse_apply_adagrad_v2(
+            var.handle, accum.handle, lr, epsilon, grad,
+            constant_op.constant(indices, dtypes.int32))
+        with ops.control_dependencies([adagrad_op]):
+          ret += i
+      return ret
+
+    # Run two tf.functions simultaneously to make sure there is no race
+    # condition between the two ops that caused deadlock before (b/270712679).
+    thread1 = threading.Thread(
+        target=lambda: self.evaluate(fn_disable_copy_on_read()))
+    thread2 = threading.Thread(
+        target=lambda: self.evaluate(fn_resource_sparse_apply_adagrad_v2()))
+    thread1.start()
+    thread2.start()
+    thread1.join()
+    thread2.join()
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/training/training_util.py

@@ -0,0 +1,421 @@
+# 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.
+# ==============================================================================
+"""Utility functions for training."""
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import graph_io
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import state_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.util.tf_export import tf_export
+
+# Picked a long key value to minimize the chance of collision with user defined
+# collection keys.
+GLOBAL_STEP_READ_KEY = 'global_step_read_op_cache'
+
+# TODO(drpng): remove this after legacy uses are resolved.
+write_graph = graph_io.write_graph
+
+
+@tf_export(v1=['train.global_step'])
+def global_step(sess, global_step_tensor):
+  """Small helper to get the global step.
+
+  ```python
+  # Create a variable to hold the global_step.
+  global_step_tensor = tf.Variable(10, trainable=False, name='global_step')
+  # Create a session.
+  sess = tf.compat.v1.Session()
+  # Initialize the variable
+  sess.run(global_step_tensor.initializer)
+  # Get the variable value.
+  print('global_step: %s' % tf.compat.v1.train.global_step(sess,
+  global_step_tensor))
+
+  global_step: 10
+  ```
+
+  Args:
+    sess: A TensorFlow `Session` object.
+    global_step_tensor:  `Tensor` or the `name` of the operation that contains
+      the global step.
+
+  Returns:
+    The global step value.
+  """
+  if context.executing_eagerly():
+    return int(global_step_tensor.numpy())
+  return sess.run(global_step_tensor).item()
+
+
+@tf_export(v1=['train.get_global_step'])
+def get_global_step(graph=None):
+  """Get the global step tensor.
+
+  The global step tensor must be an integer variable. We first try to find it
+  in the collection `GLOBAL_STEP`, or by name `global_step:0`.
+
+  Args:
+    graph: The graph to find the global step in. If missing, use default graph.
+
+  Returns:
+    The global step variable, or `None` if none was found.
+
+  Raises:
+    TypeError: If the global step tensor has a non-integer type, or if it is not
+      a `Variable`.
+
+  @compatibility(TF2)
+  With the deprecation of global graphs, TF no longer tracks variables in
+  collections. In other words, there are no global variables in TF2. Thus, the
+  global step functions have been removed  (`get_or_create_global_step`,
+  `create_global_step`, `get_global_step`) . You have two options for migrating:
+
+  1. Create a Keras optimizer, which generates an `iterations` variable. This
+     variable is automatically incremented when calling `apply_gradients`.
+  2. Manually create and increment a `tf.Variable`.
+
+  Below is an example of migrating away from using a global step to using a
+  Keras optimizer:
+
+  Define a dummy model and loss:
+
+  >>> def compute_loss(x):
+  ...   v = tf.Variable(3.0)
+  ...   y = x * v
+  ...   loss = x * 5 - x * v
+  ...   return loss, [v]
+
+  Before migrating:
+
+  >>> g = tf.Graph()
+  >>> with g.as_default():
+  ...   x = tf.compat.v1.placeholder(tf.float32, [])
+  ...   loss, var_list = compute_loss(x)
+  ...   global_step = tf.compat.v1.train.get_or_create_global_step()
+  ...   global_init = tf.compat.v1.global_variables_initializer()
+  ...   optimizer = tf.compat.v1.train.GradientDescentOptimizer(0.1)
+  ...   train_op = optimizer.minimize(loss, global_step, var_list)
+  >>> sess = tf.compat.v1.Session(graph=g)
+  >>> sess.run(global_init)
+  >>> print("before training:", sess.run(global_step))
+  before training: 0
+  >>> sess.run(train_op, feed_dict={x: 3})
+  >>> print("after training:", sess.run(global_step))
+  after training: 1
+
+  Using `get_global_step`:
+
+  >>> with g.as_default():
+  ...   print(sess.run(tf.compat.v1.train.get_global_step()))
+  1
+
+  Migrating to a Keras optimizer:
+
+  >>> optimizer = tf.keras.optimizers.SGD(.01)
+  >>> print("before training:", optimizer.iterations.numpy())
+  before training: 0
+  >>> with tf.GradientTape() as tape:
+  ...   loss, var_list = compute_loss(3)
+  ...   grads = tape.gradient(loss, var_list)
+  ...   optimizer.apply_gradients(zip(grads, var_list))
+  >>> print("after training:", optimizer.iterations.numpy())
+  after training: 1
+
+  @end_compatibility
+  """
+  graph = graph or ops.get_default_graph()
+  global_step_tensor = None
+  global_step_tensors = graph.get_collection(ops.GraphKeys.GLOBAL_STEP)
+  if len(global_step_tensors) == 1:
+    global_step_tensor = global_step_tensors[0]
+  elif not global_step_tensors:
+    try:
+      global_step_tensor = graph.get_tensor_by_name('global_step:0')
+    except KeyError:
+      return None
+  else:
+    logging.error('Multiple tensors in global_step collection.')
+    return None
+
+  assert_global_step(global_step_tensor)
+  return global_step_tensor
+
+
+@tf_export(v1=['train.create_global_step'])
+def create_global_step(graph=None):
+  """Create global step tensor in graph.
+
+  Args:
+    graph: The graph in which to create the global step tensor. If missing, use
+      default graph.
+
+  Returns:
+    Global step tensor.
+
+  Raises:
+    ValueError: if global step tensor is already defined.
+
+  @compatibility(TF2)
+  With the deprecation of global graphs, TF no longer tracks variables in
+  collections. In other words, there are no global variables in TF2. Thus, the
+  global step functions have been removed  (`get_or_create_global_step`,
+  `create_global_step`, `get_global_step`) . You have two options for migrating:
+
+  1. Create a Keras optimizer, which generates an `iterations` variable. This
+     variable is automatically incremented when calling `apply_gradients`.
+  2. Manually create and increment a `tf.Variable`.
+
+  Below is an example of migrating away from using a global step to using a
+  Keras optimizer:
+
+  Define a dummy model and loss:
+
+  >>> def compute_loss(x):
+  ...   v = tf.Variable(3.0)
+  ...   y = x * v
+  ...   loss = x * 5 - x * v
+  ...   return loss, [v]
+
+  Before migrating:
+
+  >>> g = tf.Graph()
+  >>> with g.as_default():
+  ...   x = tf.compat.v1.placeholder(tf.float32, [])
+  ...   loss, var_list = compute_loss(x)
+  ...   global_step = tf.compat.v1.train.create_global_step()
+  ...   global_init = tf.compat.v1.global_variables_initializer()
+  ...   optimizer = tf.compat.v1.train.GradientDescentOptimizer(0.1)
+  ...   train_op = optimizer.minimize(loss, global_step, var_list)
+  >>> sess = tf.compat.v1.Session(graph=g)
+  >>> sess.run(global_init)
+  >>> print("before training:", sess.run(global_step))
+  before training: 0
+  >>> sess.run(train_op, feed_dict={x: 3})
+  >>> print("after training:", sess.run(global_step))
+  after training: 1
+
+  Migrating to a Keras optimizer:
+
+  >>> optimizer = tf.keras.optimizers.SGD(.01)
+  >>> print("before training:", optimizer.iterations.numpy())
+  before training: 0
+  >>> with tf.GradientTape() as tape:
+  ...   loss, var_list = compute_loss(3)
+  ...   grads = tape.gradient(loss, var_list)
+  ...   optimizer.apply_gradients(zip(grads, var_list))
+  >>> print("after training:", optimizer.iterations.numpy())
+  after training: 1
+
+  @end_compatibility
+  """
+  graph = graph or ops.get_default_graph()
+  if get_global_step(graph) is not None:
+    raise ValueError('"global_step" already exists.')
+  if context.executing_eagerly():
+    with ops.device('cpu:0'):
+      return variable_scope.get_variable(
+          ops.GraphKeys.GLOBAL_STEP,
+          shape=[],
+          dtype=dtypes.int64,
+          initializer=init_ops.zeros_initializer(),
+          trainable=False,
+          aggregation=variables.VariableAggregation.ONLY_FIRST_REPLICA,
+          collections=[
+              ops.GraphKeys.GLOBAL_VARIABLES, ops.GraphKeys.GLOBAL_STEP
+          ])
+  # Create in proper graph and base name_scope.
+  with graph.as_default() as g, g.name_scope(None):
+    return variable_scope.get_variable(
+        ops.GraphKeys.GLOBAL_STEP,
+        shape=[],
+        dtype=dtypes.int64,
+        initializer=init_ops.zeros_initializer(),
+        trainable=False,
+        aggregation=variables.VariableAggregation.ONLY_FIRST_REPLICA,
+        collections=[ops.GraphKeys.GLOBAL_VARIABLES, ops.GraphKeys.GLOBAL_STEP])
+
+
+@tf_export(v1=['train.get_or_create_global_step'])
+def get_or_create_global_step(graph=None):
+  """Returns and create (if necessary) the global step tensor.
+
+  Args:
+    graph: The graph in which to create the global step tensor. If missing, use
+      default graph.
+
+  Returns:
+    The global step tensor.
+
+  @compatibility(TF2)
+  With the deprecation of global graphs, TF no longer tracks variables in
+  collections. In other words, there are no global variables in TF2. Thus, the
+  global step functions have been removed  (`get_or_create_global_step`,
+  `create_global_step`, `get_global_step`) . You have two options for migrating:
+
+  1. Create a Keras optimizer, which generates an `iterations` variable. This
+     variable is automatically incremented when calling `apply_gradients`.
+  2. Manually create and increment a `tf.Variable`.
+
+  Below is an example of migrating away from using a global step to using a
+  Keras optimizer:
+
+  Define a dummy model and loss:
+
+  >>> def compute_loss(x):
+  ...   v = tf.Variable(3.0)
+  ...   y = x * v
+  ...   loss = x * 5 - x * v
+  ...   return loss, [v]
+
+  Before migrating:
+
+  >>> g = tf.Graph()
+  >>> with g.as_default():
+  ...   x = tf.compat.v1.placeholder(tf.float32, [])
+  ...   loss, var_list = compute_loss(x)
+  ...   global_step = tf.compat.v1.train.get_or_create_global_step()
+  ...   global_init = tf.compat.v1.global_variables_initializer()
+  ...   optimizer = tf.compat.v1.train.GradientDescentOptimizer(0.1)
+  ...   train_op = optimizer.minimize(loss, global_step, var_list)
+  >>> sess = tf.compat.v1.Session(graph=g)
+  >>> sess.run(global_init)
+  >>> print("before training:", sess.run(global_step))
+  before training: 0
+  >>> sess.run(train_op, feed_dict={x: 3})
+  >>> print("after training:", sess.run(global_step))
+  after training: 1
+
+  Migrating to a Keras optimizer:
+
+  >>> optimizer = tf.keras.optimizers.SGD(.01)
+  >>> print("before training:", optimizer.iterations.numpy())
+  before training: 0
+  >>> with tf.GradientTape() as tape:
+  ...   loss, var_list = compute_loss(3)
+  ...   grads = tape.gradient(loss, var_list)
+  ...   optimizer.apply_gradients(zip(grads, var_list))
+  >>> print("after training:", optimizer.iterations.numpy())
+  after training: 1
+
+  @end_compatibility
+  """
+  graph = graph or ops.get_default_graph()
+  global_step_tensor = get_global_step(graph)
+  if global_step_tensor is None:
+    global_step_tensor = create_global_step(graph)
+  return global_step_tensor
+
+
+@tf_export(v1=['train.assert_global_step'])
+def assert_global_step(global_step_tensor):
+  """Asserts `global_step_tensor` is a scalar int `Variable` or `Tensor`.
+
+  Args:
+    global_step_tensor: `Tensor` to test.
+  """
+  if not (isinstance(global_step_tensor, variables.Variable) or
+          isinstance(global_step_tensor, tensor.Tensor) or
+          resource_variable_ops.is_resource_variable(global_step_tensor)):
+    raise TypeError('Existing "global_step" must be a Variable or Tensor: %s.' %
+                    global_step_tensor)
+
+  if not global_step_tensor.dtype.base_dtype.is_integer:
+    raise TypeError('Existing "global_step" does not have integer type: %s' %
+                    global_step_tensor.dtype)
+
+  if (global_step_tensor.get_shape().ndims != 0 and
+      global_step_tensor.get_shape().is_fully_defined()):
+    raise TypeError('Existing "global_step" is not scalar: %s' %
+                    global_step_tensor.get_shape())
+
+
+def _get_global_step_read(graph=None):
+  """Gets global step read tensor in graph.
+
+  Args:
+    graph: The graph in which to create the global step read tensor. If missing,
+      use default graph.
+
+  Returns:
+    Global step read tensor.
+
+  Raises:
+    RuntimeError: if multiple items found in collection GLOBAL_STEP_READ_KEY.
+  """
+  graph = graph or ops.get_default_graph()
+  global_step_read_tensors = graph.get_collection(GLOBAL_STEP_READ_KEY)
+  if len(global_step_read_tensors) > 1:
+    raise RuntimeError('There are multiple items in collection {}. '
+                       'There should be only one.'.format(GLOBAL_STEP_READ_KEY))
+
+  if len(global_step_read_tensors) == 1:
+    return global_step_read_tensors[0]
+  return None
+
+
+def _get_or_create_global_step_read(graph=None):
+  """Gets or creates global step read tensor in graph.
+
+  Args:
+    graph: The graph in which to create the global step read tensor. If missing,
+      use default graph.
+
+  Returns:
+    Global step read tensor if there is global_step_tensor else return None.
+  """
+  graph = graph or ops.get_default_graph()
+  global_step_read_tensor = _get_global_step_read(graph)
+  if global_step_read_tensor is not None:
+    return global_step_read_tensor
+  global_step_tensor = get_global_step(graph)
+  if global_step_tensor is None:
+    return None
+  # add 'zero' so that it will create a copy of variable as Tensor.
+  with graph.as_default() as g, g.name_scope(None):
+    with g.name_scope(global_step_tensor.op.name + '/'):
+      # must ensure that global_step is initialized before this run.
+      if isinstance(global_step_tensor, variables.Variable):
+        global_step_value = cond.cond(
+            variable_v1.is_variable_initialized(global_step_tensor),
+            global_step_tensor.read_value,
+            lambda: global_step_tensor.initial_value)
+      else:
+        global_step_value = global_step_tensor
+
+      global_step_read_tensor = global_step_value + 0
+      ops.add_to_collection(GLOBAL_STEP_READ_KEY, global_step_read_tensor)
+  return _get_global_step_read(graph)
+
+
+def _increment_global_step(increment, graph=None):
+  graph = graph or ops.get_default_graph()
+  global_step_tensor = get_global_step(graph)
+  if global_step_tensor is None:
+    raise ValueError(
+        'Global step tensor should be created by '
+        'tf.train.get_or_create_global_step before calling increment.')
+  global_step_read_tensor = _get_or_create_global_step_read(graph)
+  with graph.as_default() as g, g.name_scope(None):
+    with g.name_scope(global_step_tensor.op.name + '/'):
+      with ops.control_dependencies([global_step_read_tensor]):
+        return state_ops.assign_add(global_step_tensor, increment)

tensorflow/python/training/training_util_test.py

@@ -0,0 +1,122 @@
+# 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.
+# ==============================================================================
+"""Tests for training_util."""
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.platform import test
+from tensorflow.python.training import monitored_session
+from tensorflow.python.training import training_util
+
+
+class GlobalStepTest(test.TestCase):
+
+  def _assert_global_step(self, global_step, expected_dtype=dtypes.int64):
+    self.assertEqual('%s:0' % ops.GraphKeys.GLOBAL_STEP, global_step.name)
+    self.assertEqual(expected_dtype, global_step.dtype.base_dtype)
+    self.assertEqual([], global_step.get_shape().as_list())
+
+  def test_invalid_dtype(self):
+    with ops.Graph().as_default() as g:
+      self.assertIsNone(training_util.get_global_step())
+      variable_v1.VariableV1(
+          0.0,
+          trainable=False,
+          dtype=dtypes.float32,
+          name=ops.GraphKeys.GLOBAL_STEP,
+          collections=[ops.GraphKeys.GLOBAL_STEP])
+      self.assertRaisesRegex(TypeError, 'does not have integer type',
+                             training_util.get_global_step)
+    self.assertRaisesRegex(TypeError, 'does not have integer type',
+                           training_util.get_global_step, g)
+
+  def test_invalid_shape(self):
+    with ops.Graph().as_default() as g:
+      self.assertIsNone(training_util.get_global_step())
+      variable_v1.VariableV1([0],
+                             trainable=False,
+                             dtype=dtypes.int32,
+                             name=ops.GraphKeys.GLOBAL_STEP,
+                             collections=[ops.GraphKeys.GLOBAL_STEP])
+      self.assertRaisesRegex(TypeError, 'not scalar',
+                             training_util.get_global_step)
+    self.assertRaisesRegex(TypeError, 'not scalar',
+                           training_util.get_global_step, g)
+
+  def test_create_global_step(self):
+    self.assertIsNone(training_util.get_global_step())
+    with ops.Graph().as_default() as g:
+      global_step = training_util.create_global_step()
+      self._assert_global_step(global_step)
+      self.assertRaisesRegex(ValueError, 'already exists',
+                             training_util.create_global_step)
+      self.assertRaisesRegex(ValueError, 'already exists',
+                             training_util.create_global_step, g)
+      self._assert_global_step(training_util.create_global_step(ops.Graph()))
+
+  def test_get_global_step(self):
+    with ops.Graph().as_default() as g:
+      self.assertIsNone(training_util.get_global_step())
+      variable_v1.VariableV1(
+          0,
+          trainable=False,
+          dtype=dtypes.int32,
+          name=ops.GraphKeys.GLOBAL_STEP,
+          collections=[ops.GraphKeys.GLOBAL_STEP])
+      self._assert_global_step(
+          training_util.get_global_step(), expected_dtype=dtypes.int32)
+    self._assert_global_step(
+        training_util.get_global_step(g), expected_dtype=dtypes.int32)
+
+  def test_get_or_create_global_step(self):
+    with ops.Graph().as_default() as g:
+      self.assertIsNone(training_util.get_global_step())
+      self._assert_global_step(training_util.get_or_create_global_step())
+      self._assert_global_step(training_util.get_or_create_global_step(g))
+
+
+class GlobalStepReadTest(test.TestCase):
+
+  def test_global_step_read_is_none_if_there_is_no_global_step(self):
+    with ops.Graph().as_default():
+      self.assertIsNone(training_util._get_or_create_global_step_read())
+      training_util.create_global_step()
+      self.assertIsNotNone(training_util._get_or_create_global_step_read())
+
+  def test_reads_from_cache(self):
+    with ops.Graph().as_default():
+      training_util.create_global_step()
+      first = training_util._get_or_create_global_step_read()
+      second = training_util._get_or_create_global_step_read()
+      self.assertEqual(first, second)
+
+  def test_reads_before_increments(self):
+    with ops.Graph().as_default():
+      training_util.create_global_step()
+      read_tensor = training_util._get_or_create_global_step_read()
+      inc_op = training_util._increment_global_step(1)
+      inc_three_op = training_util._increment_global_step(3)
+      with monitored_session.MonitoredTrainingSession() as sess:
+        read_value, _ = sess.run([read_tensor, inc_op])
+        self.assertEqual(0, read_value)
+        read_value, _ = sess.run([read_tensor, inc_three_op])
+        self.assertEqual(1, read_value)
+        read_value = sess.run(read_tensor)
+        self.assertEqual(4, read_value)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/training/warm_starting_util.py

@@ -0,0 +1,561 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Utilities to warm-start TF.Learn Estimators."""
+
+import collections
+
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import state_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variables as variables_lib
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.training import checkpoint_ops
+from tensorflow.python.training import checkpoint_utils
+from tensorflow.python.training import saver as saver_lib
+from tensorflow.python.training.saving import saveable_object_util
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(v1=["train.VocabInfo"])
+class VocabInfo(
+    collections.namedtuple("VocabInfo", [
+        "new_vocab",
+        "new_vocab_size",
+        "num_oov_buckets",
+        "old_vocab",
+        "old_vocab_size",
+        "backup_initializer",
+        "axis",
+    ])):
+  """Vocabulary information for warm-starting.
+
+  See `tf.estimator.WarmStartSettings` for examples of using
+  VocabInfo to warm-start.
+
+  Args:
+    new_vocab: [Required] A path to the new vocabulary file (used with the model
+      to be trained).
+    new_vocab_size: [Required] An integer indicating how many entries of the new
+      vocabulary will used in training.
+    num_oov_buckets: [Required] An integer indicating how many OOV buckets are
+      associated with the vocabulary.
+    old_vocab: [Required] A path to the old vocabulary file (used with the
+      checkpoint to be warm-started from).
+    old_vocab_size: [Optional] An integer indicating how many entries of the old
+      vocabulary were used in the creation of the checkpoint. If not provided,
+      the entire old vocabulary will be used.
+    backup_initializer: [Optional] A variable initializer used for variables
+      corresponding to new vocabulary entries and OOV. If not provided, these
+      entries will be zero-initialized.
+    axis: [Optional] Denotes what axis the vocabulary corresponds to.  The
+      default, 0, corresponds to the most common use case (embeddings or
+      linear weights for binary classification / regression).  An axis of 1
+      could be used for warm-starting output layers with class vocabularies.
+
+  Returns:
+    A `VocabInfo` which represents the vocabulary information for warm-starting.
+
+  Raises:
+    ValueError: `axis` is neither 0 or 1.
+
+      Example Usage:
+```python
+      embeddings_vocab_info = tf.VocabInfo(
+          new_vocab='embeddings_vocab',
+          new_vocab_size=100,
+          num_oov_buckets=1,
+          old_vocab='pretrained_embeddings_vocab',
+          old_vocab_size=10000,
+          backup_initializer=tf.compat.v1.truncated_normal_initializer(
+              mean=0.0, stddev=(1 / math.sqrt(embedding_dim))),
+          axis=0)
+
+      softmax_output_layer_kernel_vocab_info = tf.VocabInfo(
+          new_vocab='class_vocab',
+          new_vocab_size=5,
+          num_oov_buckets=0,  # No OOV for classes.
+          old_vocab='old_class_vocab',
+          old_vocab_size=8,
+          backup_initializer=tf.compat.v1.glorot_uniform_initializer(),
+          axis=1)
+
+      softmax_output_layer_bias_vocab_info = tf.VocabInfo(
+          new_vocab='class_vocab',
+          new_vocab_size=5,
+          num_oov_buckets=0,  # No OOV for classes.
+          old_vocab='old_class_vocab',
+          old_vocab_size=8,
+          backup_initializer=tf.compat.v1.zeros_initializer(),
+          axis=0)
+
+      #Currently, only axis=0 and axis=1 are supported.
+  ```
+  """
+
+  def __new__(cls,
+              new_vocab,
+              new_vocab_size,
+              num_oov_buckets,
+              old_vocab,
+              old_vocab_size=-1,
+              backup_initializer=None,
+              axis=0):
+    if axis != 0 and axis != 1:
+      raise ValueError("The only supported values for the axis argument are 0 "
+                       "and 1.  Provided axis: {}".format(axis))
+
+    return super(VocabInfo, cls).__new__(
+        cls,
+        new_vocab,
+        new_vocab_size,
+        num_oov_buckets,
+        old_vocab,
+        old_vocab_size,
+        backup_initializer,
+        axis,
+    )
+
+
+def _infer_var_name(var):
+  """Returns name of the `var`.
+
+  Args:
+    var: A list. The list can contain either of the following:
+      (i) A single `Variable`
+      (ii) A single `ResourceVariable`
+      (iii) Multiple `Variable` objects which must be slices of the same larger
+        variable.
+      (iv) A single `PartitionedVariable`
+
+  Returns:
+    Name of the `var`
+  """
+  name_to_var_dict = saveable_object_util.op_list_to_dict(var)
+  if len(name_to_var_dict) > 1:
+    raise TypeError("`var` = %s passed as arg violates the constraints.  "
+                    "name_to_var_dict = %s" % (var, name_to_var_dict))
+  return list(name_to_var_dict.keys())[0]
+
+
+def _get_var_info(var, prev_tensor_name=None):
+  """Helper method for standarizing Variable and naming.
+
+  Args:
+    var: Current graph's variable that needs to be warm-started (initialized).
+      Can be either of the following: (i) `Variable` (ii) `ResourceVariable`
+      (iii) list of `Variable`: The list must contain slices of the same larger
+        variable. (iv) `PartitionedVariable`
+    prev_tensor_name: Name of the tensor to lookup in provided `prev_ckpt`. If
+      None, we lookup tensor with same name as given `var`.
+
+  Returns:
+    A tuple of the Tensor name and var.
+  """
+  if checkpoint_utils._is_variable(var):  # pylint: disable=protected-access
+    current_var_name = _infer_var_name([var])
+  elif (isinstance(var, list) and
+        all(checkpoint_utils._is_variable(v) for v in var)):  # pylint: disable=protected-access
+    current_var_name = _infer_var_name(var)
+  elif isinstance(var, variables_lib.PartitionedVariable):
+    current_var_name = _infer_var_name([var])
+    var = var._get_variable_list()  # pylint: disable=protected-access
+  else:
+    raise TypeError(
+        "var MUST be one of the following: a Variable, list of Variable or "
+        "PartitionedVariable, but is {}".format(type(var)))
+  if not prev_tensor_name:
+    # Assume tensor name remains the same.
+    prev_tensor_name = current_var_name
+
+  return prev_tensor_name, var
+
+
+# pylint: disable=protected-access
+# Accesses protected members of tf.Variable to reset the variable's internal
+# state.
+def _warm_start_var_with_vocab(var,
+                               current_vocab_path,
+                               current_vocab_size,
+                               prev_ckpt,
+                               prev_vocab_path,
+                               previous_vocab_size=-1,
+                               current_oov_buckets=0,
+                               prev_tensor_name=None,
+                               initializer=None,
+                               axis=0):
+  """Warm-starts given variable from `prev_tensor_name` tensor in `prev_ckpt`.
+
+  Use this method when the `var` is backed by vocabulary. This method stitches
+  the given `var` such that values corresponding to individual features in the
+  vocabulary remain consistent irrespective of changing order of the features
+  between old and new vocabularies.
+
+  Args:
+    var: Current graph's variable that needs to be warm-started (initialized).
+      Can be either of the following:
+      (i) `Variable`
+      (ii) `ResourceVariable`
+      (iii) list of `Variable`: The list must contain slices of the same larger
+        variable.
+      (iv) `PartitionedVariable`
+    current_vocab_path: Path to the vocab file used for the given `var`.
+    current_vocab_size: An `int` specifying the number of entries in the current
+      vocab.
+    prev_ckpt: A string specifying the directory with checkpoint file(s) or path
+      to checkpoint. The given checkpoint must have tensor with name
+      `prev_tensor_name` (if not None) or tensor with name same as given `var`.
+    prev_vocab_path: Path to the vocab file used for the tensor in `prev_ckpt`.
+    previous_vocab_size: If provided, will constrain previous vocab to the first
+      `previous_vocab_size` entries.  -1 means use the entire previous vocab.
+    current_oov_buckets: An `int` specifying the number of out-of-vocabulary
+      buckets used for given `var`.
+    prev_tensor_name: Name of the tensor to lookup in provided `prev_ckpt`. If
+      None, we lookup tensor with same name as given `var`.
+    initializer: Variable initializer to be used for missing entries.  If None,
+      missing entries will be zero-initialized.
+    axis: Axis of the variable that the provided vocabulary corresponds to.
+
+  Raises:
+    ValueError: If required args are not provided.
+  """
+  if not (current_vocab_path and current_vocab_size and prev_ckpt and
+          prev_vocab_path):
+    raise ValueError("Invalid args: Must provide all of [current_vocab_path, "
+                     "current_vocab_size, prev_ckpt, prev_vocab_path}.")
+  if checkpoint_utils._is_variable(var):
+    var = [var]
+  elif (isinstance(var, list) and
+        all(checkpoint_utils._is_variable(v) for v in var)):
+    var = var
+  elif isinstance(var, variables_lib.PartitionedVariable):
+    var = var._get_variable_list()
+  else:
+    raise TypeError(
+        "var MUST be one of the following: a Variable, list of Variable or "
+        "PartitionedVariable, but is {}".format(type(var)))
+
+  if not prev_tensor_name:
+    # Assume tensor name remains the same.
+    prev_tensor_name = _infer_var_name(var)
+
+  total_v_first_axis = sum(v.get_shape().as_list()[0] for v in var)
+  for v in var:
+    v_shape = v.get_shape().as_list()
+    slice_info = v._get_save_slice_info()
+    partition_info = None
+    if slice_info:
+      partition_info = variable_scope._PartitionInfo(
+          full_shape=slice_info.full_shape, var_offset=slice_info.var_offset)
+
+    if axis == 0:
+      new_row_vocab_size = current_vocab_size
+      new_col_vocab_size = v_shape[1]
+      old_row_vocab_size = previous_vocab_size
+      old_row_vocab_file = prev_vocab_path
+      new_row_vocab_file = current_vocab_path
+      old_col_vocab_file = None
+      new_col_vocab_file = None
+      num_row_oov_buckets = current_oov_buckets
+      num_col_oov_buckets = 0
+    elif axis == 1:
+      # Note that we must compute this value across all partitions, whereas
+      # in the axis = 0 case, we can simply use v_shape[1] because we don't
+      # allow partitioning across axis = 1.
+      new_row_vocab_size = total_v_first_axis
+      new_col_vocab_size = current_vocab_size
+      old_row_vocab_size = -1
+      old_row_vocab_file = None
+      new_row_vocab_file = None
+      old_col_vocab_file = prev_vocab_path
+      new_col_vocab_file = current_vocab_path
+      num_row_oov_buckets = 0
+      num_col_oov_buckets = current_oov_buckets
+    else:
+      raise ValueError("The only supported values for the axis argument are 0 "
+                       "and 1.  Provided axis: {}".format(axis))
+
+    init = checkpoint_ops._load_and_remap_matrix_initializer(
+        ckpt_path=checkpoint_utils._get_checkpoint_filename(prev_ckpt),
+        old_tensor_name=prev_tensor_name,
+        new_row_vocab_size=new_row_vocab_size,
+        new_col_vocab_size=new_col_vocab_size,
+        old_row_vocab_size=old_row_vocab_size,
+        old_row_vocab_file=old_row_vocab_file,
+        new_row_vocab_file=new_row_vocab_file,
+        old_col_vocab_file=old_col_vocab_file,
+        new_col_vocab_file=new_col_vocab_file,
+        num_row_oov_buckets=num_row_oov_buckets,
+        num_col_oov_buckets=num_col_oov_buckets,
+        initializer=initializer)
+    new_init_val = ops.convert_to_tensor(
+        init(shape=v_shape, partition_info=partition_info))
+    v._initializer_op = state_ops.assign(v, new_init_val)
+
+
+# pylint: enable=protected-access
+
+
+def _get_grouped_variables(vars_to_warm_start):
+  """Collects and groups (possibly partitioned) variables into a dictionary.
+
+  The variables can be provided explicitly through vars_to_warm_start, or they
+  are retrieved from collections (see below).
+
+  Args:
+    vars_to_warm_start: One of the following:
+
+      - A regular expression (string) that captures which variables to
+        warm-start (see tf.compat.v1.get_collection).  This expression will
+        only consider variables in the TRAINABLE_VARIABLES collection.
+      - A list of strings, each representing a full variable name to warm-start.
+        These will consider variables in GLOBAL_VARIABLES collection.
+      - A list of Variables to warm-start.
+      - `None`, in which case all variables in TRAINABLE_VARIABLES will be used.
+  Returns:
+    A dictionary mapping variable names (strings) to lists of Variables.
+  Raises:
+    ValueError: If vars_to_warm_start is not a string, `None`, a list of
+      `Variables`, or a list of strings.
+  """
+  # TODO(b/143899805): Remove unicode checks when deprecating Python2.
+  if isinstance(vars_to_warm_start, str) or vars_to_warm_start is None:
+    # Both vars_to_warm_start = '.*' and vars_to_warm_start = None will match
+    # everything (in TRAINABLE_VARIABLES) here.
+    logging.info("Warm-starting variables only in TRAINABLE_VARIABLES.")
+    list_of_vars = ops.get_collection(
+        ops.GraphKeys.TRAINABLE_VARIABLES, scope=vars_to_warm_start)
+  elif isinstance(vars_to_warm_start, list):
+    if all(isinstance(v, str) for v in vars_to_warm_start):
+      list_of_vars = []
+      for v in vars_to_warm_start:
+        list_of_vars += ops.get_collection(
+            ops.GraphKeys.GLOBAL_VARIABLES, scope=v)
+    elif all(checkpoint_utils._is_variable(v) for v in vars_to_warm_start):  # pylint: disable=protected-access
+      list_of_vars = vars_to_warm_start
+    else:
+      raise ValueError("If `vars_to_warm_start` is a list, it must be all "
+                       "`Variable` or all `str`.  Given types are {}".format(
+                           [type(v) for v in vars_to_warm_start]))
+  else:
+    raise ValueError("`vars_to_warm_start must be a `list` or `str`.  Given "
+                     "type is {}".format(type(vars_to_warm_start)))
+  # We have to deal with partitioned variables, since get_collection flattens
+  # out the list.
+  grouped_variables = {}
+  for v in list_of_vars:
+    t = [v] if not isinstance(v, list) else v
+    var_name = _infer_var_name(t)
+    grouped_variables.setdefault(var_name, []).append(v)
+
+  return grouped_variables
+
+
+def _get_object_checkpoint_renames(path, variable_names):
+  """Returns a dictionary mapping variable names to checkpoint keys.
+
+  The warm-starting utility expects variable names to match with the variable
+  names in the checkpoint. For object-based checkpoints, the variable names
+  and names in the checkpoint are different. Thus, for object-based checkpoints,
+  this function is used to obtain the map from variable names to checkpoint
+  keys.
+
+  Args:
+    path: path to checkpoint directory or file.
+    variable_names: list of variable names to load from the checkpoint.
+
+  Returns:
+    If the checkpoint is object-based, this function returns a map from variable
+    names to their corresponding checkpoint keys.
+    If the checkpoint is name-based, this returns an empty dict.
+
+  Raises:
+    ValueError: If the object-based checkpoint is missing variables.
+  """
+  fname = checkpoint_utils._get_checkpoint_filename(path)  # pylint: disable=protected-access
+  try:
+    names_to_keys = saver_lib.object_graph_key_mapping(fname)
+  except errors.NotFoundError:
+    # If an error is raised from `object_graph_key_mapping`, then the
+    # checkpoint is name-based. There are no renames, so return an empty dict.
+    return {}
+
+  missing_names = set(variable_names) - set(names_to_keys.keys())
+  if missing_names:
+    raise ValueError(
+        "Attempting to warm-start from an object-based checkpoint, but found "
+        "that the checkpoint did not contain values for all variables. The "
+        "following variables were missing: {}"
+        .format(missing_names))
+  return {name: names_to_keys[name] for name in variable_names}
+
+
+@tf_export(v1=["train.warm_start"])
+def warm_start(ckpt_to_initialize_from,
+               vars_to_warm_start=".*",
+               var_name_to_vocab_info=None,
+               var_name_to_prev_var_name=None):
+  """Warm-starts a model using the given settings.
+
+  If you are using a tf.estimator.Estimator, this will automatically be called
+  during training.
+
+  Args:
+    ckpt_to_initialize_from: [Required] A string specifying the directory with
+      checkpoint file(s) or path to checkpoint from which to warm-start the
+      model parameters.
+    vars_to_warm_start: [Optional] One of the following:
+
+      - A regular expression (string) that captures which variables to
+        warm-start (see tf.compat.v1.get_collection).  This expression will only
+        consider variables in the TRAINABLE_VARIABLES collection -- if you need
+        to warm-start non_TRAINABLE vars (such as optimizer accumulators or
+        batch norm statistics), please use the below option.
+      - A list of strings, each a regex scope provided to
+        tf.compat.v1.get_collection with GLOBAL_VARIABLES (please see
+        tf.compat.v1.get_collection).  For backwards compatibility reasons,
+        this is separate from the single-string argument type.
+      - A list of Variables to warm-start.  If you do not have access to the
+        `Variable` objects at the call site, please use the above option.
+      - `None`, in which case only TRAINABLE variables specified in
+        `var_name_to_vocab_info` will be warm-started.
+
+      Defaults to `'.*'`, which warm-starts all variables in the
+      TRAINABLE_VARIABLES collection.  Note that this excludes variables such
+      as accumulators and moving statistics from batch norm.
+    var_name_to_vocab_info: [Optional] Dict of variable names (strings) to
+      `tf.estimator.VocabInfo`. The variable names should be "full" variables,
+      not the names of the partitions.  If not explicitly provided, the variable
+      is assumed to have no (changes to) vocabulary.
+    var_name_to_prev_var_name: [Optional] Dict of variable names (strings) to
+      name of the previously-trained variable in `ckpt_to_initialize_from`. If
+      not explicitly provided, the name of the variable is assumed to be same
+      between previous checkpoint and current model.  Note that this has no
+      effect on the set of variables that is warm-started, and only controls
+      name mapping (use `vars_to_warm_start` for controlling what variables to
+      warm-start).
+
+  Raises:
+    ValueError: If the WarmStartSettings contains prev_var_name or VocabInfo
+      configuration for variable names that are not used.  This is to ensure
+      a stronger check for variable configuration than relying on users to
+      examine the logs.
+  """
+  logging.info("Warm-starting from: {}".format(ckpt_to_initialize_from))
+  grouped_variables = _get_grouped_variables(vars_to_warm_start)
+
+  if var_name_to_vocab_info is None:
+    var_name_to_vocab_info = {}
+
+  if not var_name_to_prev_var_name:
+    # Detect whether the checkpoint is object-based, in which case the
+    # var_name_to_prev_var_name dictionary should map variable names to
+    # checkpoint keys. If the user has specified var_name_to_prev_var_name, we
+    # do not override it.
+    var_name_to_prev_var_name = _get_object_checkpoint_renames(
+        ckpt_to_initialize_from, grouped_variables.keys())
+
+  warmstarted_count = 0
+
+  # Keep track of which var_names in var_name_to_prev_var_name and
+  # var_name_to_vocab_info have been used.  Err on the safer side by throwing an
+  # exception if any are unused by the end of the loop.  It is easy to misname
+  # a variable during this configuration, in which case without this check, we
+  # would fail to warm-start silently.
+  prev_var_name_used = set()
+  vocab_info_used = set()
+
+  # Group the vocabless vars into one call to init_from_checkpoint.
+  vocabless_vars = {}
+  for var_name, variable in grouped_variables.items():
+    prev_var_name = var_name_to_prev_var_name.get(var_name)
+    if prev_var_name:
+      prev_var_name_used.add(var_name)
+    vocab_info = var_name_to_vocab_info.get(var_name)
+    if vocab_info:
+      vocab_info_used.add(var_name)
+      warmstarted_count += 1
+      logging.debug(
+          "Warm-starting variable: {}; current_vocab: {} current_vocab_size: {}"
+          " prev_vocab: {} prev_vocab_size: {} current_oov: {} prev_tensor: {}"
+          " initializer: {}".format(
+              var_name, vocab_info.new_vocab, vocab_info.new_vocab_size,
+              vocab_info.old_vocab, (vocab_info.old_vocab_size if
+                                     vocab_info.old_vocab_size > 0 else "All"),
+              vocab_info.num_oov_buckets, prev_var_name or "Unchanged",
+              vocab_info.backup_initializer or "zero-initialized"))
+      _warm_start_var_with_vocab(
+          variable,
+          current_vocab_path=vocab_info.new_vocab,
+          current_vocab_size=vocab_info.new_vocab_size,
+          prev_ckpt=ckpt_to_initialize_from,
+          prev_vocab_path=vocab_info.old_vocab,
+          previous_vocab_size=vocab_info.old_vocab_size,
+          current_oov_buckets=vocab_info.num_oov_buckets,
+          prev_tensor_name=prev_var_name,
+          initializer=vocab_info.backup_initializer,
+          axis=vocab_info.axis)
+    else:
+      # For the special value of vars_to_warm_start = None,
+      # we only warm-start variables with explicitly specified vocabularies.
+      if vars_to_warm_start:
+        warmstarted_count += 1
+        logging.debug("Warm-starting variable: {}; prev_var_name: {}".format(
+            var_name, prev_var_name or "Unchanged"))
+        # Because we use a default empty list in grouped_variables, single
+        # unpartitioned variables will be lists here, which we rectify in order
+        # for init_from_checkpoint logic to work correctly.
+        if len(variable) == 1:
+          variable = variable[0]
+        prev_tensor_name, var = _get_var_info(variable, prev_var_name)
+        if prev_tensor_name in vocabless_vars:
+          # The API for checkpoint_utils.init_from_checkpoint accepts a mapping
+          # from checkpoint tensor names to model variable names, so it does not
+          # support warm-starting two variables from the same tensor.  Our work-
+          # around is to run init_from_checkpoint multiple times, each time we
+          # encounter a new variable that should be initialized by a previously-
+          # used tensor.
+          logging.debug("Requested prev_var_name {} initialize both {} and {}; "
+                        "calling init_from_checkpoint.".format(
+                            prev_tensor_name,
+                            vocabless_vars[prev_tensor_name],
+                            var))
+          checkpoint_utils.init_from_checkpoint(ckpt_to_initialize_from,
+                                                vocabless_vars)
+          vocabless_vars.clear()
+        vocabless_vars[prev_tensor_name] = var
+
+  if vocabless_vars:
+    checkpoint_utils.init_from_checkpoint(ckpt_to_initialize_from,
+                                          vocabless_vars)
+  prev_var_name_not_used = set(
+      var_name_to_prev_var_name.keys()) - prev_var_name_used
+  vocab_info_not_used = set(var_name_to_vocab_info.keys()) - vocab_info_used
+
+  logging.info("Warm-started %d variables.", warmstarted_count)
+
+  if prev_var_name_not_used:
+    raise ValueError(
+        "You provided the following variables in "
+        "var_name_to_prev_var_name that were not used: "
+        "{0}.  Perhaps you misspelled them?  Here is the list of viable "
+        "variable names: {1}".format(prev_var_name_not_used,
+                                     grouped_variables.keys()))
+  if vocab_info_not_used:
+    raise ValueError(
+        "You provided the following variables in "
+        "var_name_to_vocab_info that were not used: {0}. "
+        " Perhaps you misspelled them?  Here is the list of viable variable "
+        "names: {1}".format(vocab_info_not_used, grouped_variables.keys()))