def _defun_with_scope(self, scope):
    """Creates a defun wrapped inside a variable creator scope."""

    weak_wrapped_fn = None
    def wrapped_fn(*args, **kwds):
      """Wraps `self._python_function` in a variable creator scope."""
      # We register a variable creator with reduced priority. If an outer
      # variable creator is just modifying keyword arguments to the variable
      # constructor, this will work harmoniously. Since the `scope` registered
      # here actually creates the variable, it taking priority would otherwise
      # ignore the outer creator.
      #
      # If an outer variable creator calls the variable constructor manually,
      # for example creating a MirroredVariable, then they won't call our
      # creator. This means we won't be able to trace the initialization graph,
      # and so variable initializers can't depend on function arguments. This is
      # better than the alternative, tracing the initialization graph but giving
      # the user a variable type they didn't want.
      with ops.get_default_graph()._variable_creator_scope(scope, priority=50):  # pylint: disable=protected-access
        # __wrapped__ allows AutoGraph to swap in a converted function. We give
        # the function a weak reference to itself to avoid a reference cycle.
        return weak_wrapped_fn().__wrapped__(*args, **kwds)
    weak_wrapped_fn = weakref.ref(wrapped_fn)

    # TODO(mdan): Pipe self._experimental_autograph_options through.
    return function_lib.defun(
        tf_decorator.make_decorator(self._python_function, wrapped_fn),
        input_signature=self._input_signature,
        autograph=self._autograph,
        experimental_autograph_options=self._experimental_autograph_options)

 def testBasic(self):
   matmul = function.defun(math_ops.matmul)
   t = constant_op.constant([[1.0, 2.0], [3.0, 4.0]])
   sq = matmul(t, t, transpose_a=True)
   sq2 = matmul(sq, t, transpose_a=True)
   self.assertAllEqual(sq.numpy().reshape(-1), [10, 14, 14, 20])
   self.assertAllEqual(sq2.numpy().reshape(-1), [52, 76, 74, 108])

  def __init__(self, name=None, use_global_variables=False):
    self._built = False
    self._vars = []
    self._initial_values = {}
    self._updates = []
    self._use_global_variables = use_global_variables
    name = name or self.__class__.__name__
    # Replace things like spaces in name to create a valid scope name.
    scope_name = _to_replace.sub("_", name)
    # We create the variable scope now to get the unique name that will
    # be used as a variable prefix when build() calls add_variable().
    with variable_scope.variable_scope(
        scope_name, use_resource=True, reuse=False) as scope:
      pos = scope.name.rfind(scope_name)
      self._name = name + scope.name[pos + len(scope_name):]
      self._scope = scope

    # Ensures that if the user calls build directly we still set self._built to
    # True to prevent variables from being recreated.
    self._build = self.build

    def actual_build(*args, **kwargs):
      self._build(*args, **kwargs)
      self._built = True
    self.build = actual_build
    self.build.__doc__ = self._build.__doc__

    # Captures construction scope for proper initialization.
    if context.executing_eagerly():
      self._construction_scope = context.eager_mode
    else:
      # We make self.call() into a graph callable here, so that we can
      # return a single op that performs all of the variable updates.
      self._construction_scope = ops.get_default_graph().as_default
      self.call = function.defun(self.call)

def _defun_with_scope(scope, fn, input_signature):

  def wrapped_fn(*args, **kwds):
    with variable_scope.variable_creator_scope(scope):
      return fn(*args, **kwds)

  return function_lib.defun(wrapped_fn, input_signature=input_signature)

  def __call__(self, *args, **kwds):
    """Calls the graph function."""
    if self._created_variables:
      # In this case we have created variables on the first call, so we run the
      # defunned version which is guaranteed to never create variables.
      return self._stateless_fn(*args, **kwds)  # pylint: disable=not-callable
    elif self._stateful_fn is not None:
      # In this case we have not created variables on the first call. So we can
      # run the first trace but we should fail if variables are created.
      results = self._first_trace(*args, **kwds)
      if self._created_variables:
        raise ValueError("Creating variables on a non-first call to a function"
                         " decorated with tf.function.")
      return results

    self._initialize(args, kwds)

    if not self._created_variables:
      # If we did not create any variables the trace we have is good enough.
      return _call_concrete(self._concrete_stateful_fn, args, kwds)()

    def fn_with_cond(*inner_args, **inner_kwds):
      """Conditionally runs initialization if it's needed."""
      condition = True
      for variable in self._created_variables:
        condition = condition and resource_variable_ops.var_is_initialized_op(
            variable.handle)
      # We want to call stateless_fn if possible because it avoids recomputing
      # potentially expensive initializers.
      return control_flow_ops.cond(
          condition,
          lambda: self._stateless_fn(*inner_args, **inner_kwds),
          _call_concrete(self._concrete_stateful_fn, inner_args, inner_kwds))

    return function_lib.defun(fn_with_cond)(*args, **kwds)

def _defun_with_scope(scope, fn):

  def wrapped_fn(*args, **kwds):
    with variable_scope.variable_creator_scope(scope):
      return fn(*args, **kwds)

  return function.defun(wrapped_fn)

 def testSequenceInputs(self):
   clip_by_global_norm = function.defun(clip_ops.clip_by_global_norm)
   t_list = [tensor.Tensor(1.0), tensor.Tensor(2.0)]
   clipped_list, global_norm = clip_by_global_norm(t_list, tensor.Tensor(.2))
   for t in clipped_list:
     self.assertTrue(isinstance(t, tensor.Tensor))
   self.assertTrue(isinstance(global_norm, tensor.Tensor))

  def testFunctionOnDevice(self):
    if not context.context().num_gpus():
      self.skipTest('No GPUs found')

    x = constant_op.constant([1.]).gpu()
    f = function.defun(math_ops.add)
    y = f(x, x).cpu()
    self.assertAllEqual(y, [2.])

 def _benchmark_defun_matmul(self,
                             m,
                             transpose_b,
                             num_iters,
                             execution_mode=None):
   f = function.defun(math_ops.matmul)
   func = lambda: f(m, m, transpose_b=transpose_b)
   self._run(func, num_iters, execution_mode=execution_mode)

 def _defun(self, fn):
   """Returns a defun generated from the input function."""
   # TODO(mdan): Pipe self._experimental_autograph_options through.
   return function_lib.defun(
       fn,
       input_signature=self.input_signature,
       autograph=self._autograph,
       experimental_autograph_options=self._experimental_autograph_options)

  def testFunctionOnDevice(self):
    if not context.context().num_gpus():
      self.skipTest('No GPUs found')

    x = tensor.Tensor([1.]).as_gpu_tensor()
    f = function.defun(math_ops.add)
    y = f(x, x).as_cpu_tensor()
    self.assertAllEqual(y.numpy(), [2.])

  def testGradient(self):
    matmul = function.defun(math_ops.matmul)

    def sq(x):
      return matmul(x, x, transpose_a=True)

    t = constant_op.constant([[1.0, 2.0], [3.0, 4.0]])
    grad_t, = backprop.gradients_function(sq, [0])(t)
    self.assertAllEqual(grad_t, [[6, 6], [14, 14]])

  def testGatherResourceWithDefun(self):
    with ops.device('cpu:0'):
      v = resource_variable_ops.ResourceVariable([0.0, 1.0, 2.0])

    def sum_gather():
      return math_ops.reduce_sum(array_ops.gather(v, [1, 2]))

    defined = function.defun(sum_gather)
    self.assertAllEqual(sum_gather(), defined())

  def testFunctionHandlesInputsPlacedOnTheWrongDeviceGracefully(self):
    if not context.context().num_gpus():
      self.skipTest('No GPUs found')

    # The Reshape op requires the shape tensor to be placed in host memory.
    reshape = function.defun(array_ops.reshape)
    value = constant_op.constant([1., 2.])
    shape = constant_op.constant([2, 1]).gpu()
    reshape(value, shape)  # No error is raised

  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, `func` will be executed 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 `name` is None.
    """
    if create_graph_function:
      self._func = function.defun(func)
    else:
      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 to
    # completion, which is not the same as whether the scope has been created.
    self._variables_created = False
    # `MirroredStrategy` builds the graph with multiple threads. If a
    # `merge_call` happens within a template, multiple calls may be in progress
    # simultaneously. This variable keeps track of whether any call of the
    # template has started.
    self._first_call = True

  def testFunctionHandlesInputsOnDifferentDevices(self):
    if not context.context().num_gpus():
      self.skipTest('No GPUs found')

    # The Reshape op requires the shape tensor to be placed in host memory.
    reshape = function.defun(array_ops.reshape)
    value = constant_op.constant([1., 2.]).gpu()
    shape = constant_op.constant([2, 1])
    reshaped = reshape(value, shape).cpu()
    self.assertAllEqual(reshaped, [[1], [2]])

  def testEagerPyFuncInDefun(self):
    with test_util.device(use_gpu=True):
      def wrapper():
        a = array_ops.ones((3, 3), dtype=dtypes.float32)
        x = array_ops.ones((3, 1), dtype=dtypes.float32)
        return script_ops.eager_py_func(matmul, inp=[a, x], Tout=dtypes.float32)

      wrapped = function.defun(wrapper)
      ret = self.evaluate(wrapped())
      self.assertAllEqual(ret, [[3.0], [3.0], [3.0]])

  def testBasicGraphMode(self):
    matmul = function.defun(math_ops.matmul)

    @function.defun
    def sq(a):
      return matmul(a, a)

    t = constant_op.constant([[1.0, 2.0], [3.0, 4.0]])
    out = sq(t)
    self.assertAllEqual(out, math_ops.matmul(t, t).numpy())

  def testDefunShapeInferenceWithCapturedResourceVariable(self):
    v = resource_variable_ops.ResourceVariable([[1, 2], [3, 4]])

    def f():
      x = constant_op.constant([[1, 2], [3, 4]])
      out = math_ops.matmul(v, x)
      self.assertEqual(out.get_shape(), tensor_shape.TensorShape([2, 2]))

    compiled = function.defun(f)
    compiled()

  def benchmark_defun_without_signature(self):

    def func(t1, t2, t3, t4, t5, t6, t7, t8):
      del t1, t2, t3, t4, t5, t6, t7, t8
      return None

    defined = function.defun(func)
    t = constant_op.constant(0.0)
    cache_computation = lambda: defined(t, t, t, t, t, t, t, t)
    self._run(cache_computation, 30000)