text: """Saver for eager mode TensorFlow."""
  # 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.
  # ==============================================================================
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import contextlib
  
  from tensorflow.python.eager import context
  from tensorflow.python.framework import errors
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import resource_variable_ops
  from tensorflow.python.training import checkpoint_utils
  from tensorflow.python.training import saver as _saver
  
  
  def _init_from_checkpoint(self, *args, **kwargs):
    """Overrides default init by loading value from checkpoint."""
    # pylint: disable=protected-access
    self._old_init(*args, **kwargs)
    ckpt_name = self._map_func(self._shared_name)
    if ckpt_name not in self._ckpt_var_cache:
      raise errors.NotFoundError(None, None,
                                 "%s not found in checkpoint" % ckpt_name)
  
    val = self._ckpt_var_cache.get(ckpt_name, None)
    if val is not None:
      self.assign(val)
      # Avoid assigning for the second time.
      self._ckpt_var_cache[ckpt_name] = None
    # pylint: enable=protected-access
  
  
  @contextlib.contextmanager
  def restore_variables_on_create(save_path, map_func=None):
    """ContextManager that restores variables on creation.
  
      When save_path is None (e.g. No checkpoint), does nothing.
      Otherwise, it preloads all values from checkpoint. When the
      corresponding variable is first created, it assigns the checkpoint
      value to the variable.
  
      ```python
      with restore_variables_on_create(
          tf.train.latest_checkpoint(checkpoint_dir)):
      ```
  
    Args:
      save_path: The checkpoint file prefix.
      map_func: A function that given the variable name as argument
          and returns a variable name in checkpoint for restore. If
          None, use the variable with the same name in checkpoint to restore.
          It's an error that the mapped variable name doesn't exist in
          checkpoint.
  
    Yields:
      Nothing.
  
    Raises:
      NotFoundError: If the variable is not found in checkpoint.
      ValueError: If not used in eager mode or map_func is not callable.
    """
    if context.in_graph_mode():
      raise ValueError(
          "Currently, restore_variables_on_create can only be used with "
          "eager execution enabled.")
    if save_path:
      if map_func is None:
        map_func_wrapper = lambda self, x: x
      else:
        if not callable(map_func):
          raise ValueError("map_func must be callaled.")
        map_func_wrapper = lambda self, x: map_func(x)
  
      ckpt_var_cache = dict()
      reader = checkpoint_utils.load_checkpoint(save_path)
      for k, _ in checkpoint_utils.list_variables(save_path):
        ckpt_var_cache[k] = reader.get_tensor(k)
  
      old_init = getattr(resource_variable_ops.ResourceVariable,
                         "_init_from_args", None)
      assert old_init, "ResourceVariable misses _init_from_args method."
      setattr(resource_variable_ops.ResourceVariable, "_init_from_args",
              _init_from_checkpoint)
      setattr(resource_variable_ops.ResourceVariable, "_old_init", old_init)
      setattr(resource_variable_ops.ResourceVariable, "_map_func",
              map_func_wrapper)
      setattr(resource_variable_ops.ResourceVariable, "_ckpt_var_cache",
              ckpt_var_cache)
    try:
      yield
    except Exception as e:
      raise e
    finally:
      if save_path:
        setattr(resource_variable_ops.ResourceVariable, "_init_from_args",
                old_init)
        setattr(resource_variable_ops.ResourceVariable, "_old_init", None)
        setattr(resource_variable_ops.ResourceVariable, "_map_func", None)
        setattr(resource_variable_ops.ResourceVariable, "_ckpt_var_cache", None)
  
  
  class Saver(object):
    """A tf.train.Saver adapter for use when eager execution is enabled.
    """
  
    def __init__(self, var_list):
      """A  tf.train.Saver adapter for use when eager execution is enabled.
  
        The API, and on-disk format, mimic tf.train.Saver except that no
        Session is needed.
  
      Args:
        var_list: The list of variables that will be saved and restored. Either a
          list of `tfe.Variable` objects, or a dictionary mapping names to
          `tfe.Variable` objects.
  
      Raises:
        RuntimeError: if invoked when eager execution has not been enabled.
      """
      if context.in_graph_mode():
        raise RuntimeError("tfe.Saver can only be used when eager "
                           "execution is enabled. Use tf.train.Saver when "
                           "building graphs.")
      self._saver = _saver.Saver(var_list=var_list)
  
    def save(self, file_prefix, global_step=None):
      """Saves variables.
  
      Args:
        file_prefix: Path prefix of files created for the checkpoint.
        global_step: If provided the global step number is appended to file_prefix
          to create the checkpoint filename. The optional argument can be a
          Tensor, a Variable, or an integer.
  
      Returns:
        A string: prefix of filenames created for the checkpoint. This may be
         an extension of file_prefix that is suitable to pass as an argument
         to a subsequent call to `restore()`.
      """
      with ops.device("/device:CPU:0"):
        return self._saver.save(
            None, file_prefix, write_meta_graph=False, global_step=global_step)
  
    def restore(self, file_prefix):
      """Restores previously saved variables.
  
      Args:
        file_prefix: Path prefix where parameters were previously saved.
          Typically obtained from a previous `save()` call, or from
          @{tf.train.latest_checkpoint}.
      """
      with ops.device("/device:CPU:0"):
        self._saver.restore(None, file_prefix)
  
  
  def get_optimizer_variables(optimizer):
    """Returns a list of variables for the given `tf.train.Optimizer`.
  
    Equivalent to `optimizer.variables()`.
  
    Args:
      optimizer: An instance of `tf.train.Optimizer` which has created variables
        (typically after a call to `Optimizer.minimize`).
    Returns:
      A list of variables which have been created by the `Optimizer`.
    """
    return optimizer.variables()
  