text: # Copyright 2015 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  
  """Provides templates which allow variable sharing."""
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import functools
  import traceback
  
  from tensorflow.python.eager import context
  from tensorflow.python.eager import function
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import variable_scope
  from tensorflow.python.platform import tf_logging as logging
  from tensorflow.python.util import tf_contextlib
  from tensorflow.python.util import tf_decorator
  from tensorflow.python.util.deprecation import deprecated
  
  
  __all__ = ["make_template"]
  
  
  def make_template(name_, func_, create_scope_now_=False, unique_name_=None,
                    custom_getter_=None, **kwargs):
    """Given an arbitrary function, wrap it so that it does variable sharing.
  
    This wraps `func_` in a Template and partially evaluates it. Templates are
    functions that create variables the first time they are called and reuse them
    thereafter. In order for `func_` to be compatible with a `Template` it must
    have the following properties:
  
    * The function should create all trainable variables and any variables that
       should be reused by calling `tf.get_variable`. If a trainable variable is
       created using `tf.Variable`, then a ValueError will be thrown. Variables
       that are intended to be locals can be created by specifying
       `tf.Variable(..., trainable=false)`.
    * The function may use variable scopes and other templates internally to
        create and reuse variables, but it shouldn't use `tf.global_variables` to
        capture variables that are defined outside of the scope of the function.
    * Internal scopes and variable names should not depend on any arguments that
        are not supplied to `make_template`. In general you will get a ValueError
        telling you that you are trying to reuse a variable that doesn't exist
        if you make a mistake.
  
    In the following example, both `z` and `w` will be scaled by the same `y`. It
    is important to note that if we didn't assign `scalar_name` and used a
    different name for z and w that a `ValueError` would be thrown because it
    couldn't reuse the variable.
  
    ```python
    def my_op(x, scalar_name):
      var1 = tf.get_variable(scalar_name,
                             shape=[],
                             initializer=tf.constant_initializer(1))
      return x * var1
  
    scale_by_y = tf.make_template('scale_by_y', my_op, scalar_name='y')
  
    z = scale_by_y(input1)
    w = scale_by_y(input2)
    ```
  
    As a safe-guard, the returned function will raise a `ValueError` after the
    first call if trainable variables are created by calling `tf.Variable`.
  
    If all of these are true, then 2 properties are enforced by the template:
  
    1. Calling the same template multiple times will share all non-local
        variables.
    2. Two different templates are guaranteed to be unique, unless you reenter the
        same variable scope as the initial definition of a template and redefine
        it. An examples of this exception:
  
    ```python
    def my_op(x, scalar_name):
      var1 = tf.get_variable(scalar_name,
                             shape=[],
                             initializer=tf.constant_initializer(1))
      return x * var1
  
    with tf.variable_scope('scope') as vs:
      scale_by_y = tf.make_template('scale_by_y', my_op, scalar_name='y')
      z = scale_by_y(input1)
      w = scale_by_y(input2)
  
    # Creates a template that reuses the variables above.
    with tf.variable_scope(vs, reuse=True):
      scale_by_y2 = tf.make_template('scale_by_y', my_op, scalar_name='y')
      z2 = scale_by_y2(input1)
      w2 = scale_by_y2(input2)
    ```
  
    Depending on the value of `create_scope_now_`, the full variable scope may be
    captured either at the time of first call or at the time of construction. If
    this option is set to True, then all Tensors created by repeated calls to the
    template will have an extra trailing _N+1 to their name, as the first time the
    scope is entered in the Template constructor no Tensors are created.
  
    Note: `name_`, `func_` and `create_scope_now_` have a trailing underscore to
    reduce the likelihood of collisions with kwargs.
  
    Args:
      name_: A name for the scope created by this template. If necessary, the name
        will be made unique by appending `_N` to the name.
      func_: The function to wrap.
      create_scope_now_: Boolean controlling whether the scope should be created
        when the template is constructed or when the template is called. Default
        is False, meaning the scope is created when the template is called.
      unique_name_: When used, it overrides name_ and is not made unique. If a
        template of the same scope/unique_name already exists and reuse is false,
        an error is raised. Defaults to None.
      custom_getter_: Optional custom getter for variables used in `func_`. See
        the @{tf.get_variable} `custom_getter` documentation for
        more information.
      **kwargs: Keyword arguments to apply to `func_`.
  
    Returns:
      A function to encapsulate a set of variables which should be created once
      and reused. An enclosing scope will be created either when `make_template`
      is called or when the result is called, depending on the value of
      `create_scope_now_`. Regardless of the value, the first time the template
      is called it will enter the scope with no reuse, and call `func_` to create
      variables, which are guaranteed to be unique. All subsequent calls will
      re-enter the scope and reuse those variables.
  
    Raises:
      ValueError: if the name is None.
    """
    return make_template_internal(
        name_,
        func_,
        create_scope_now_,
        unique_name_,
        custom_getter_,
        create_graph_function_=False,
        **kwargs)
  
  
  def make_template_internal(name_,
                             func_,
                             create_scope_now_=False,
                             unique_name_=None,
                             custom_getter_=None,
                             create_graph_function_=False,
                             **kwargs):
    """Make a template, optionally compiling func_ into a graph function.
  
    See `make_template` for full documentation.
  
    Args:
      name_: A name for the scope created by this template. If necessary, the name
        will be made unique by appending `_N` to the name.
      func_: The function to wrap.
      create_scope_now_: Boolean controlling whether the scope should be created
        when the template is constructed or when the template is called. Default
        is False, meaning the scope is created when the template is called.
      unique_name_: When used, it overrides name_ and is not made unique. If a
        template of the same scope/unique_name already exists and reuse is false,
        an error is raised. Defaults to None. If executing eagerly, must be None.
      custom_getter_: Optional custom getter for variables used in `func_`. See
        the @{tf.get_variable} `custom_getter` documentation for
        more information.
      create_graph_function_: When True, the first invocation of the template will
        execute `func_` as is, to allow for variable creation; however, the second
        invocation and every invocation thereafter will execute func as a graph
        function.  In particular, this implies that `func_` must satisfy the
        properties that `function.defun` requires of functions: See the
        documentation of `function.defun` for details. When executing eagerly,
        setting this flag to True can improve performance. Regardless of whether
        eager execution is enabled, enabling this flag gives the caller access to
        graph-function semantics, i.e., accesses to variables are totally ordered
        and side-effecting ops are not pruned.
      **kwargs: Keyword arguments to apply to `func_`.
  
    Returns:
      A function to encapsulate a set of variables which should be created once
      and reused. An enclosing scope will be created either when `make_template`
      is called or when the result is called, depending on the value of
      `create_scope_now_`. Regardless of the value, the first time the template
      is called it will enter the scope with no reuse, and call `func_` to create
      variables, which are guaranteed to be unique. All subsequent calls will
      re-enter the scope and reuse those variables.
  
    Raises:
      ValueError: if the name is None.
      ValueError: if unique_name_ is not None and eager execution is enabled.
    """
  
    if kwargs:
      func_ = tf_decorator.make_decorator(func_, functools.partial(
          func_, **kwargs))
    if context.in_eager_mode():
      if unique_name_ is not None:
        raise ValueError(
            "unique_name_ cannot be used when eager exeuction is enabled.")
      return EagerTemplate(
          name_,
          func_,
          create_scope_now=create_scope_now_,
          custom_getter=custom_getter_,
          create_graph_function=create_graph_function_)
    return Template(
        name_,
        func_,
        create_scope_now=create_scope_now_,
        unique_name=unique_name_,
        custom_getter=custom_getter_,
        create_graph_function=create_graph_function_)
  
  
  def _skip_common_stack_elements(stacktrace, base_case):
    """Skips items that the target stacktrace shares with the base stacktrace."""
    for i, (trace, base) in enumerate(zip(stacktrace, base_case)):
      if trace != base:
        return stacktrace[i:]
    return stacktrace[-1:]
  
  
  class Template(object):
    """Wrap a function to aid in variable sharing.
  
    Templates are functions that create variables the first time they are called
    and reuse them thereafter. See `make_template` for full documentation.
  
    Note: By default, the full variable scope is captured at the time of first
    call. If `create_scope_now_` is passed as True to the constructor, the full
    scope will be captured there, but no variables will created until the first
    call.
    """
  
    def __init__(self, name, func, create_scope_now=False, unique_name=None,
                 custom_getter=None, create_graph_function=False):
      """Creates a template for the given function.
  
      Args:
        name: A name for the scope created by this template. The
          name will be made unique by appending `_N` to the it (see how
          `tf.variable_scope` treats the `default_name` for details).
        func: The function to apply each time.
        create_scope_now: Whether to create the scope at Template construction
          time, rather than first call. Defaults to false. Creating the scope at
          construction time may be more convenient if the template is to passed
          through much lower level code, and you want to be sure of the scope
          name without knowing exactly where it will be first called. If set to
          True, the scope will be created in the constructor, and all subsequent
          times in `__call__`, leading to a trailing numeral being added to the
          names of all created Tensors. If set to False, the scope will be created
          at the first call location.
        unique_name: When used, it overrides `name` and is not made unique. If a
          template of the same scope/unique_name already exists and reuse is
          false, an error is raised. Defaults to None.
        custom_getter: optional custom getter to pass to `variable_scope()`
        create_graph_function: When True, the first invocation of the template
          will execute `func` as is, to allow for variable creation; however, the
          second invocation and every invocation thereafter will execute `func` as
          a graph function. Enabling this flag gives the caller access to
          graph-function semantics, i.e., accesses to variables are totally
          ordered and side-effecting ops are not pruned.
  
  
      Raises:
        ValueError: if the name is None.
      """
      self._func = func
      self._stacktrace = traceback.format_stack()[:-2]
      self._name = name
      self._unique_name = unique_name
      self._custom_getter = custom_getter
      if name is None:
        raise ValueError("name cannot be None.")
      if create_scope_now:
        with variable_scope._pure_variable_scope(  # pylint:disable=protected-access
            (self._unique_name or
             variable_scope._get_unique_variable_scope(self._name)),  # pylint:disable=protected-access
            custom_getter=self._custom_getter) as vs:
          self._variable_scope = vs
      else:
        self._variable_scope = None
      # This variable keeps track of whether the template has been called yet,
      # which is not the same as whether the scope has been created.
      self._variables_created = False
      self._create_graph_function = create_graph_function
  
    def _call_func(self, args, kwargs):
      try:
        vars_at_start = len(ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES))
        trainable_at_start = len(
            ops.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES))
  
        result = self._func(*args, **kwargs)
        if self._create_graph_function and not self._variables_created:
          # Only execute self._func as a graph function once variables are
          # created.
          self._func = function.defun(self._func)
  
        if self._variables_created:
          # Variables were previously created, implying this is not the first
          # time the template has been called. Check to make sure that no new
          # trainable variables were created this time around.
          trainable_variables = ops.get_collection(
              ops.GraphKeys.TRAINABLE_VARIABLES)
          # If a variable that we intend to train is created as a side effect
          # of creating a template, then that is almost certainly an error.
          if trainable_at_start != len(trainable_variables):
            raise ValueError("Trainable variable created when calling a template "
                             "after the first time, perhaps you used tf.Variable "
                             "when you meant tf.get_variable: %s" %
                             (trainable_variables[trainable_at_start:],))
  
          # Non-trainable tracking variables are a legitimate reason why a new
          # variable would be created, but it is a relatively advanced use-case,
          # so log it.
          variables = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)
          if vars_at_start != len(variables):
            logging.info("New variables created when calling a template after "
                         "the first time, perhaps you used tf.Variable when you "
                         "meant tf.get_variable: %s",
                         variables[vars_at_start:])
        else:
          self._variables_created = True
        return result
      except Exception as exc:
        # Reraise the exception, but append the original definition to the
        # trace.
        args = exc.args
        if not args:
          arg0 = ""
        else:
          arg0 = args[0]
        trace = "".join(_skip_common_stack_elements(self._stacktrace,
                                                    traceback.format_stack()))
        arg0 = "%s\n\noriginally defined at:\n%s" % (arg0, trace)
        new_args = [arg0]
        new_args.extend(args[1:])
        exc.args = tuple(new_args)
        raise
  
    def __call__(self, *args, **kwargs):
      if self._variable_scope:
        # Only reuse variables if they were already created.
        with variable_scope.variable_scope(
            self._variable_scope, reuse=self._variables_created):
          result = self._call_func(args, kwargs)
        return result
      else:
        # The scope was not created at construction time, so create it here.
        # Subsequent calls should reuse variables.
        with variable_scope.variable_scope(
            self._unique_name, self._name,
            custom_getter=self._custom_getter) as vs:
          self._variable_scope = vs
          result = self._call_func(args, kwargs)
          return result
  
    @property
    def name(self):
      """Returns the name given to this Template."""
      return self._name
  
    @property
    def func(self):
      """Returns the func given to this Template."""
      return self._func
  
    @property
    def variable_scope(self):
      """Returns the variable scope object created by this Template."""
      return self._variable_scope
  
    @property
    def variable_scope_name(self):
      """Returns the variable scope name created by this Template."""
      if self._variable_scope:
        name = self._variable_scope.name
        # To prevent partial matches on the scope_name, we add '/' at the end.
        return name if name[-1] == "/" else name + "/"
  
    @property
    def variables(self):
      """Returns the list of global and local variables created by the Template.
      """
      return self.global_variables + self.local_variables
  
    @property
    def trainable_variables(self):
      """Returns the list of trainable variables created by the Template."""
      if self._variables_created:
        return ops.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES,
                                  self.variable_scope_name)
      else:
        return []
  
    @property
    def non_trainable_variables(self):
      """Returns the list of non-trainable variables created by the Template."""
      # TODO(apassos) Make sure it matches Eager when using local variables.
      global_variables = self.global_variables
      trainable_variables = set(self.trainable_variables)
      return [x for x in global_variables if x not in trainable_variables]
  
    @property
    def global_variables(self):
      """Returns the list of global variables created by the Template."""
      if self._variables_created:
        return ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES,
                                  self.variable_scope_name)
      else:
        return []
  
    @property
    def local_variables(self):
      """Returns the list of global variables created by the Template."""
      if self._variables_created:
        return ops.get_collection(ops.GraphKeys.LOCAL_VARIABLES,
                                  self.variable_scope_name)
      else:
        return []
  
    @property
    def weights(self):
      """List of weights/variables created by the Template."""
      return self.variables
  
    @property
    def trainable_weights(self):
      """List of trainable weights/variables created by the Template."""
      return self.trainable_variables
  
    @property
    def non_trainable_weights(self):
      """List of non-trainable weights/variables created by the Template."""
      return self.non_trainable_variables
  
    @property
    @deprecated(
        "2017-02-21", "The .var_scope property is deprecated. Please change your "
        "code to use the .variable_scope property")
    def var_scope(self):
      """Returns the variable scope object created by this Template."""
      return self._variable_scope
  
  
  class _EagerTemplateVariableStore(object):
    """Wrapper around EagerVariableStore to support nesting EagerTemplates.
    """
  
    def __init__(self, variable_scope_name):
      self._variable_scope_name = variable_scope_name
      default = variable_scope._get_default_variable_store()  # pylint: disable=protected-access
      if default._store_eager_variables:  # pylint: disable=protected-access
        self._eager_variable_store = variable_scope.EagerVariableStore(default)
      else:
        self._eager_variable_store = variable_scope.EagerVariableStore()
  
    def set_variable_scope_name(self, variable_scope_name):
      self._variable_scope_name = variable_scope_name
  
    @tf_contextlib.contextmanager
    def as_default(self):
      try:
        with self._eager_variable_store.as_default():
          yield
      finally:
        # Each _EagerTemplateVariableStore object lives underneath a variable
        # scope (see EagerTemplate.__call__). This variable scope's subscopes are
        # closed when the EagerTemplate object returns from __call__. For
        # top-level _EagerTemplateVariableStore objects, the variable store to
        # which the variable scope is attached is different from the
        # EagerVariableStore; as such it is necessary to close its subscopes
        # here as well.
        if self._variable_scope_name is None:
          raise RuntimeError("A variable scope must be set before an "
                             "_EagerTemplateVariableStore object exits.")
        self._eager_variable_store._store.close_variable_subscopes(  # pylint: disable=protected-access
            self._variable_scope_name)
  
    def _variables_in_scope(self, variable_list):
      if self._variable_scope_name is None:
        raise RuntimeError(
            "A variable scope must be set before variables can be accessed.")
      return [
          v for v in variable_list
          if v.name.startswith(self._variable_scope_name + "/")
      ]
  
    def variables(self):
      return self._variables_in_scope(self._eager_variable_store.variables())
  
    def trainable_variables(self):
      return self._variables_in_scope(
          self._eager_variable_store.trainable_variables())
  
    def non_trainable_variables(self):
      return self._variables_in_scope(
          self._eager_variable_store.non_trainable_variables())
  
  
  class EagerTemplate(Template):
    """Wrap a function to aid in variable sharing in Eager mode.
  
    Templates are functions that create variables the first time they are called
    and reuse them thereafter. See `make_template` for full documentation.
  
    Note: By default, the full variable scope is captured at the time of first
    call. If `create_scope_now` is passed as True to the constructor, the full
    scope will be captured there, but no variables will be created until the first
    call.
    """
  
    def __init__(self, name, func, create_scope_now=False, custom_getter=None,
                 create_graph_function=False):
      """Creates a template for the given function.
  
      Args:
        name: A name for the scope created by this template. The
          name will be made unique by appending `_N` to the it (see how
          `tf.variable_scope` treats the `default_name` for details).
        func: The function to apply each time.
        create_scope_now: Whether to create the scope at Template construction
          time, rather than first call. Defaults to false. Creating the scope at
          construction time may be more convenient if the template is passed
          through much lower level code, and you want to be sure of the scope
          name without knowing exactly where it will be first called. If set to
          True, the scope will be created in the constructor, and all subsequent
          times in `__call__`, leading to a trailing numeral being added to the
          names of all created Tensors. If set to False, the scope will be created
          at the first call location.
        custom_getter: optional custom getter to pass to `variable_scope()`
        create_graph_function: When True, the first invocation of the template
          will execute `func` as is, to allow for variable creation; however, the
          second invocation and every invocation thereafter will execute `func` as
          a graph function. Enabling this flag allows the caller to reap the
          performance benefits associated with executing graphs, at the cost of
          sacrificing debuggability; however, not all functions can be compiled
          into graph functions. See the documentation for `function.defun` for
          details.
  
      Raises:
        RuntimeError: if eager execution is not enabled.
      """
      if not context.in_eager_mode():
        raise RuntimeError(
            "{} objects can only be used when eager execution is enabled, use "
            "tf.Template for graph construction".
            format(type(self)))
      super(EagerTemplate, self).__init__(name, func, create_scope_now, None,
                                          custom_getter, create_graph_function)
      if self._variable_scope is not None:
        variable_scope_name = self._variable_scope.name
      else:
        # Defer setting the variable scope name until the variable scope
        # is created in __call__.
        variable_scope_name = None
      self._template_store = _EagerTemplateVariableStore(variable_scope_name)
  
    def _call_func(self, args, kwargs):
      try:
        vars_at_start = self._template_store.variables()
        trainable_at_start = self._template_store.trainable_variables()
  
        result = self._func(*args, **kwargs)
        if self._create_graph_function and not self._variables_created:
          # Only execute self._func as a graph function once variables are
          # created.
          self._func = function.defun(self._func)
  
        if self._variables_created:
          # Variables were previously created, implying this is not the first
          # time the template has been called. Check to make sure that no new
          # trainable variables were created this time around.
          trainable_variables = self._template_store.trainable_variables()
          # If a variable that we intend to train is created as a side effect
          # of creating a template, then that is almost certainly an error.
          if len(trainable_at_start) != len(trainable_variables):
            raise ValueError("Trainable variable created when calling a template "
                             "after the first time, perhaps you used tf.Variable "
                             "when you meant tf.get_variable: %s" %
                             list(set(trainable_variables) -
                                  set(trainable_at_start)))
  
          # Non-trainable tracking variables are a legitimate reason why a new
          # variable would be created, but it is a relatively advanced use-case,
          # so log it.
          variables = self._template_store.variables()
          if len(vars_at_start) != len(variables):
            logging.info("New variables created when calling a template after "
                         "the first time, perhaps you used tf.Variable when you "
                         "meant tf.get_variable: %s",
                         list(set(variables) - set(vars_at_start)))
        else:
          self._variables_created = True
        return result
      except Exception as exc:
        # Reraise the exception, but append the original definition to the
        # trace.
        args = exc.args
        if not args:
          arg0 = ""
        else:
          arg0 = args[0]
        trace = "".join(_skip_common_stack_elements(self._stacktrace,
                                                    traceback.format_stack()))
        arg0 = "%s\n\noriginally defined at:\n%s" % (arg0, trace)
        new_args = [arg0]
        new_args.extend(args[1:])
        exc.args = tuple(new_args)
        raise
  
    def __call__(self, *args, **kwargs):
      # In both branches below, the template store is installed as default after
      # the variable scope is opened in order to ensure that templates nested at
      # the same level correctly uniquify lower variable scope names.
      if self._variable_scope:
        with variable_scope.variable_scope(
            self._variable_scope, reuse=variable_scope.AUTO_REUSE):
          with self._template_store.as_default():
            result = self._call_func(args, kwargs)
        return result
      else:
        # The scope was not created at construction time, so create it here.
        # Subsequent calls should reuse variables.
        with variable_scope.variable_scope(
            self._unique_name, self._name,
            custom_getter=self._custom_getter) as vs:
          self._variable_scope = vs
          # Because the scope was not created at construction time, the template
          # store's variable scope name is unset; set it here.
          self._template_store.set_variable_scope_name(vs.name)
          with self._template_store.as_default():
            result = self._call_func(args, kwargs)
          return result
  
    @property
    def name(self):
      """Returns the name given to this Template."""
      return self._name
  
    @property
    def func(self):
      """Returns the func given to this Template."""
      return self._func
  
    @property
    def variable_scope(self):
      """Returns the variable scope object created by this Template."""
      return self._variable_scope
  
    @property
    def variable_scope_name(self):
      """Returns the variable scope name created by this Template."""
      if self._variable_scope:
        name = self._variable_scope.name
        # To prevent partial matches on the scope_name, we add '/' at the end.
        return name if name[-1] == "/" else name + "/"
  
    @property
    def variables(self):
      """Returns the list of variables created by the Template."""
      # Currently there is no local variable in Eager mode.
      if not self._variables_created:
        return []
      return self._template_store.variables()
  
    @property
    def trainable_variables(self):
      """Returns the list of trainable variables created by the Template."""
      # Currently there is no local variable in Eager mode.
      if not self._variables_created:
        return []
      return self._template_store.trainable_variables()
  
    @property
    def non_trainable_variables(self):
      """Returns the list of non-trainable variables created by the Template."""
      # Currently there is no local variable in Eager mode.
      if not self._variables_created:
        return []
      return self._template_store.non_trainable_variables()
  
    @property
    def global_variables(self):
      """Returns the list of global variables created by the Template."""
      # Currently there is no local variable in Eager mode.
      if not self._variables_created:
        return []
      return self.variables
  
    @property
    def local_variables(self):
      """Returns the list of global variables created by the Template."""
      # Currently there is no local variable in Eager mode.
      return []
  