# Copyright 2018 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 itertools

from tensorflow.python.eager import backprop
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
from tensorflow.python.eager import function
from tensorflow.python.framework import auto_control_deps as acd
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 control_flow_ops
from tensorflow.python.ops import gen_resource_variable_ops
from tensorflow.python.ops import gen_sendrecv_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 adam
from tensorflow.python.training import momentum


class AutomaticControlDependenciesTest(test.TestCase):

  def testBasic(self):
    with context.graph_mode(), self.cached_session():
      v = resource_variable_ops.ResourceVariable(1.0)
      self.evaluate(variables.global_variables_initializer())
      with acd.AutomaticControlDependencies() as c:
        v.assign(v + 1)
        v.assign(2 * v)
        val = v.read_value()
        val = c.mark_as_return(val)
      self.assertAllEqual(val, 4.0)

  def testNoControlDepsBetweenVariableReads(self):
    with context.graph_mode(), self.cached_session():
      v = resource_variable_ops.ResourceVariable(1.0)
      self.evaluate(variables.global_variables_initializer())
      with acd.AutomaticControlDependencies():
        read_op1 = gen_resource_variable_ops.read_variable_op(
            v.handle, v.dtype).op
        read_op2 = gen_resource_variable_ops.read_variable_op(
            v.handle, v.dtype).op
        gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
      self.assertNotIn(read_op1, read_op2.control_inputs)
      self.assertNotIn(read_op2, read_op1.control_inputs)

  def testVariableReadThenWrite(self):
    with context.graph_mode(), self.cached_session():
      v = resource_variable_ops.ResourceVariable(1.0)
      self.evaluate(variables.global_variables_initializer())
      with acd.AutomaticControlDependencies():
        read_op1 = gen_resource_variable_ops.read_variable_op(
            v.handle, v.dtype).op
        read_op2 = gen_resource_variable_ops.read_variable_op(
            v.handle, v.dtype).op
        assign_op = gen_resource_variable_ops.assign_variable_op(
            v.handle, v + 1)
      # Writes should have control deps from "all" reads since last write
      # or start of the code block.
      self.assertIn(read_op1, assign_op.control_inputs)
      self.assertIn(read_op2, assign_op.control_inputs)
      # There should be no control deps between reads.
      self.assertNotIn(read_op1, read_op2.control_inputs)
      self.assertNotIn(read_op2, read_op1.control_inputs)

  def testVariableWriteThenRead(self):
    with context.graph_mode(), self.cached_session():
      v = resource_variable_ops.ResourceVariable(1.0)
      self.evaluate(variables.global_variables_initializer())
      with acd.AutomaticControlDependencies():
        assign_op = gen_resource_variable_ops.assign_variable_op(
            v.handle, v + 1)
        read_op1 = gen_resource_variable_ops.read_variable_op(
            v.handle, v.dtype).op
        read_op2 = gen_resource_variable_ops.read_variable_op(
            v.handle, v.dtype).op
      # Reads should have a control dep from the last write.
      self.assertIn(assign_op, read_op1.control_inputs)
      self.assertIn(assign_op, read_op2.control_inputs)
      # There should be no control deps between reads.
      self.assertNotIn(read_op1, read_op2.control_inputs)
      self.assertNotIn(read_op2, read_op1.control_inputs)

  def testVariableReadsInOpsWithMustRun(self):
    with context.graph_mode(), self.cached_session():
      v = resource_variable_ops.ResourceVariable(1.0)
      self.evaluate(variables.global_variables_initializer())
      with acd.AutomaticControlDependencies() as c:
        read_op = gen_resource_variable_ops.read_variable_op(v.handle,
                                                             v.dtype).op
        # Read ops get added to control outputs only if they have consumers.
        c.mark_as_return(read_op.outputs[0])
      self.assertIn(read_op, c.ops_which_must_run)

  def testVariableMultipleReadsAndWrites(self):
    with context.graph_mode(), self.cached_session():
      v = resource_variable_ops.ResourceVariable(1.0)
      self.evaluate(variables.global_variables_initializer())
      with acd.AutomaticControlDependencies() as c:
        # 2 reads -> 2 writes -> 2 reads -> 2 writes.
        read_op1 = gen_resource_variable_ops.read_variable_op(
            v.handle, v.dtype).op
        read_op2 = gen_resource_variable_ops.read_variable_op(
            v.handle, v.dtype).op
        assign_op1 = gen_resource_variable_ops.assign_variable_op(
            v.handle, v + 1)
        assign_op2 = gen_resource_variable_ops.assign_variable_op(
            v.handle, v + 1)
        read_op3 = gen_resource_variable_ops.read_variable_op(
            v.handle, v.dtype).op
        read_op4 = gen_resource_variable_ops.read_variable_op(
            v.handle, v.dtype).op
        assign_op3 = gen_resource_variable_ops.assign_variable_op(
            v.handle, v + 1)
        assign_op4 = gen_resource_variable_ops.assign_variable_op(
            v.handle, v + 1)
        # Read ops get added to control outputs only if they have consumers.
        c.mark_as_return(read_op1.outputs[0])
        c.mark_as_return(read_op2.outputs[0])
        c.mark_as_return(read_op3.outputs[0])
        c.mark_as_return(read_op4.outputs[0])

      # Verify the control edges.
      self.assertIn(read_op1, assign_op1.control_inputs)
      self.assertIn(read_op2, assign_op1.control_inputs)
      self.assertIn(assign_op1, assign_op2.control_inputs)
      self.assertIn(assign_op2, read_op3.control_inputs)
      self.assertIn(assign_op2, read_op4.control_inputs)
      self.assertIn(read_op3, assign_op3.control_inputs)
      self.assertIn(read_op4, assign_op3.control_inputs)
      self.assertIn(assign_op3, assign_op4.control_inputs)

      # There should be no control deps between reads.
      read_ops = [read_op1, read_op2, read_op3, read_op4]
      for src_op, tgt_op in itertools.product(read_ops, read_ops):
        self.assertNotIn(src_op, tgt_op.control_inputs)

      # Reads must be in `ops_which_must_run`.
      self.assertIn(read_op1, c.ops_which_must_run)
      self.assertIn(read_op2, c.ops_which_must_run)
      self.assertIn(read_op3, c.ops_which_must_run)
      self.assertIn(read_op4, c.ops_which_must_run)
      # Last write must be in `ops_which_must_run`.
      self.assertIn(assign_op4, c.ops_which_must_run)

  def testSendInOpsWithMustRun(self):
    with context.graph_mode(), self.cached_session():
      v = resource_variable_ops.ResourceVariable(1.0)
      self.evaluate(variables.global_variables_initializer())
      with acd.AutomaticControlDependencies() as c:
        send_op = gen_sendrecv_ops.send(v, "x", "/", 0, "/")

      # Send must be in `ops_which_must_run`.
      self.assertIn(send_op, c.ops_which_must_run)

  def _testVariableReadInFunctionalOp(self, build_functional_op, op_type):
    v = resource_variable_ops.ResourceVariable(1.0)
    self.evaluate(variables.global_variables_initializer())

    @def_function.function
    def read_var_in_while():
      gen_resource_variable_ops.read_variable_op(
          v.handle, v.dtype, name="read1")

      result = build_functional_op(v)
      gen_resource_variable_ops.read_variable_op(
          v.handle, v.dtype, name="read2")
      gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
      return result

    func_graph = read_var_in_while.get_concrete_function().graph
    assert len(func_graph.inputs) == 1

    def get_op(op_type, sub_name):
      operations = [
          op for op in func_graph.get_operations()
          if op.type == op_type and sub_name in op.name
      ]
      assert len(operations) == 1
      return operations[0]

    read1 = get_op("ReadVariableOp", "read1")
    functional_op = get_op(op_type, "")
    read2 = get_op("ReadVariableOp", "read2")
    assign_op = get_op("AssignVariableOp", "")
    # Since the functional op only has reads, previous reads e.g. read1 do not\
    # have a control edge to it and next future reads e.g. read2 do not have a
    # control edge from it.
    self.assertNotIn(read1, functional_op.control_inputs)
    self.assertNotIn(functional_op, read2.control_inputs)
    self.assertIn(read1, assign_op.control_inputs)
    self.assertIn(read2, assign_op.control_inputs)
    self.assertIn(functional_op, assign_op.control_inputs)

  def testVariableReadInWhileLoop(self):

    def build_functional_op(v):

      def body(_):
        return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)

      return control_flow_ops.while_loop(
          lambda i: True, body, [0.0], maximum_iterations=1)

    self._testVariableReadInFunctionalOp(build_functional_op, "While")

  def testVariableReadInCondTrueBranch(self):

    def build_functional_op(v):

      def then_branch():
        return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)

      def else_branch():
        return array_ops.zeros([], v.dtype)

      return control_flow_ops.cond(
          constant_op.constant(True), then_branch, else_branch)

    self._testVariableReadInFunctionalOp(build_functional_op, "If")

  def testVariableReadInCondFalseBranch(self):

    def build_functional_op(v):

      def then_branch():
        return array_ops.zeros([], v.dtype)

      def else_branch():
        return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)

      return control_flow_ops.cond(
          constant_op.constant(False), then_branch, else_branch)

    self._testVariableReadInFunctionalOp(build_functional_op, "If")

  def testVariableReadInCaseBranch0(self):

    def build_functional_op(v):

      def branch0():
        return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)

      def branch1():
        return array_ops.zeros([], v.dtype)

      return control_flow_ops.switch_case(
          constant_op.constant(0), [branch0, branch1])

    self._testVariableReadInFunctionalOp(build_functional_op, "Case")

  def testVariableReadInCaseBranch1(self):

    def build_functional_op(v):

      def branch0():
        return array_ops.zeros([], v.dtype)

      def branch1():
        return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)

      return control_flow_ops.switch_case(
          constant_op.constant(0), [branch0, branch1])

    self._testVariableReadInFunctionalOp(build_functional_op, "Case")

  def testVariableReadInFunction(self):

    def build_functional_op(v):

      @def_function.function
      def fn_with_read():
        return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)

      return fn_with_read()

    self._testVariableReadInFunctionalOp(build_functional_op,
                                         "StatefulPartitionedCall")

  def testVariableReadInNestedFunction(self):

    def build_functional_op(v):

      @def_function.function
      def fn_with_read():

        @def_function.function
        def inner_fn():
          return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)

        return inner_fn()

      return fn_with_read()

    self._testVariableReadInFunctionalOp(build_functional_op,
                                         "StatefulPartitionedCall")

  def testVariableReadInWhileInInnerFunc(self):

    def build_functional_op(v):

      @def_function.function
      def fn_with_read():

        @def_function.function
        def inner_fn():

          def body(_):
            return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)

          return control_flow_ops.while_loop(
              lambda i: True, body, [0.0], maximum_iterations=1)

        return inner_fn()

      return fn_with_read()

    self._testVariableReadInFunctionalOp(build_functional_op,
                                         "StatefulPartitionedCall")

  def testVariableReadInCondInInnerFunc(self):

    def build_functional_op(v):

      @def_function.function
      def fn_with_read():

        @def_function.function
        def inner_fn():

          def then_branch():
            return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)

          def else_branch():
            return array_ops.zeros([], v.dtype)

          return control_flow_ops.cond(
              constant_op.constant(True), then_branch, else_branch)

        return inner_fn()

      return fn_with_read()

    self._testVariableReadInFunctionalOp(build_functional_op,
                                         "StatefulPartitionedCall")

  def _testVariableWriteInFunctionalOp(self, build_functional_op, op_type):
    v = resource_variable_ops.ResourceVariable(1.0)
    self.evaluate(variables.global_variables_initializer())

    @def_function.function
    def write_var_in_while():
      gen_resource_variable_ops.read_variable_op(
          v.handle, v.dtype, name="read1")

      result = build_functional_op(v)
      gen_resource_variable_ops.read_variable_op(
          v.handle, v.dtype, name="read2")
      gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
      return result

    func_graph = write_var_in_while.get_concrete_function().graph
    assert len(func_graph.inputs) == 1

    def get_op(op_type, sub_name):
      operations = [
          op for op in func_graph.get_operations()
          if op.type == op_type and sub_name in op.name
      ]
      assert len(operations) == 1
      return operations[0]

    read1 = get_op("ReadVariableOp", "read1")
    functional_op = get_op(op_type, "")
    read2 = get_op("ReadVariableOp", "read2")
    assign_op = get_op("AssignVariableOp", "")
    # Since the While has writes, it has control edges from previous reads
    # e.g. `read1` and to future reads(`read2`) and writes(`assign_op`).
    self.assertIn(read1, functional_op.control_inputs)
    self.assertIn(functional_op, read2.control_inputs)
    self.assertIn(read2, assign_op.control_inputs)
    self.assertIn(functional_op, assign_op.control_inputs)

  def testVariableWriteInWhileLoop(self):

    def build_functional_op(v):

      def body(_):
        gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
        return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)

      return control_flow_ops.while_loop(
          lambda i: True, body, [0.0], maximum_iterations=1)

    self._testVariableWriteInFunctionalOp(build_functional_op, "While")

  def testVariableWriteInCondTrueBranch(self):

    def build_functional_op(v):

      def then_branch():
        gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
        return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)

      def else_branch():
        return array_ops.zeros([], v.dtype)

      return control_flow_ops.cond(
          constant_op.constant(True), then_branch, else_branch)

    self._testVariableWriteInFunctionalOp(build_functional_op, "If")

  def testVariableWriteInCondFalseBranch(self):

    def build_functional_op(v):

      def then_branch():
        return array_ops.zeros([], v.dtype)

      def else_branch():
        gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
        return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)

      return control_flow_ops.cond(
          constant_op.constant(False), then_branch, else_branch)

    self._testVariableWriteInFunctionalOp(build_functional_op, "If")

  def testVariableWriteInCaseBranch0(self):

    def build_functional_op(v):

      def branch0():
        gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
        return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)

      def branch1():
        return array_ops.zeros([], v.dtype)

      return control_flow_ops.switch_case(
          constant_op.constant(0), [branch0, branch1])

    self._testVariableWriteInFunctionalOp(build_functional_op, "Case")

  def testVariableWriteInCaseBranch1(self):

    def build_functional_op(v):

      def branch0():
        return array_ops.zeros([], v.dtype)

      def branch1():
        gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
        return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)

      return control_flow_ops.switch_case(
          constant_op.constant(0), [branch0, branch1])

    self._testVariableWriteInFunctionalOp(build_functional_op, "Case")

  def testVariableWriteInFunction(self):

    def build_functional_op(v):

      @def_function.function
      def fn_with_write():
        gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
        return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)

      return fn_with_write()

    self._testVariableWriteInFunctionalOp(build_functional_op,
                                          "StatefulPartitionedCall")

  def testVariableWriteInNestedFunction(self):

    def build_functional_op(v):

      @def_function.function
      def fn_with_write():

        @def_function.function
        def inner_fn():
          gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
          return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)

        return inner_fn()

      return fn_with_write()

    self._testVariableWriteInFunctionalOp(build_functional_op,
                                          "StatefulPartitionedCall")

  def testVariableWriteInWhileInInnerFunc(self):

    def build_functional_op(v):

      @def_function.function
      def fn_with_write():

        @def_function.function
        def inner_fn():

          def body(_):
            gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
            return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)

          return control_flow_ops.while_loop(
              lambda i: True, body, [0.0], maximum_iterations=1)

        return inner_fn()

      return fn_with_write()

    self._testVariableWriteInFunctionalOp(build_functional_op,
                                          "StatefulPartitionedCall")

  def testVariableWriteInCondInInnerFunc(self):

    def build_functional_op(v):

      @def_function.function
      def fn_with_write():

        @def_function.function
        def inner_fn():

          def then_branch():
            gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
            return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)

          def else_branch():
            return array_ops.zeros([], v.dtype)

          return control_flow_ops.cond(
              constant_op.constant(True), then_branch, else_branch)

        return inner_fn()

      return fn_with_write()

    self._testVariableWriteInFunctionalOp(build_functional_op,
                                          "StatefulPartitionedCall")

  @test_util.run_v1_only("b/120545219")
  def testCondMustRun(self):
    with context.graph_mode(), self.cached_session():
      v = resource_variable_ops.ResourceVariable(1.0)
      self.evaluate(variables.global_variables_initializer())
      p = array_ops.placeholder(dtype=dtypes.bool)
      with acd.AutomaticControlDependencies() as c:

        def true_fn():
          v.assign(v + 1)
          return 0.0

        def false_fn():
          v.assign(v + 4)
          return 1.0

        control_flow_ops.cond(p, true_fn, false_fn)
        val = v.read_value()
        val = c.mark_as_return(val)
      self.assertAllEqual(val.eval(feed_dict={p: False}), 5.0)
      self.assertAllEqual(val.eval(feed_dict={p: True}), 6.0)

  @test_util.run_v1_only("b/120545219")
  def testCondMustRunSeparateRead(self):
    with context.graph_mode(), self.cached_session():
      v = resource_variable_ops.ResourceVariable(1.0)
      self.evaluate(variables.global_variables_initializer())
      p = array_ops.placeholder(dtype=dtypes.bool)
      with acd.AutomaticControlDependencies() as c:

        def true_fn():
          v.assign(v + 1)
          return 0.0

        def false_fn():
          v.assign(v + 4)
          return 1.0

        control_flow_ops.cond(p, true_fn, false_fn)
        one = constant_op.constant(1.0)
        one = c.mark_as_return(one)
      one.eval(feed_dict={p: False})
      self.assertAllEqual(v.read_value(), 5.0)
      one.eval(feed_dict={p: True})
      self.assertAllEqual(v.read_value(), 6.0)

  @test_util.run_v1_only("b/120545219")
  def testCondNested(self):
    with context.graph_mode(), self.cached_session():
      v = resource_variable_ops.ResourceVariable(1.0)
      self.evaluate(variables.global_variables_initializer())
      p = array_ops.placeholder(dtype=dtypes.bool)
      q = array_ops.placeholder(dtype=dtypes.bool)
      with acd.AutomaticControlDependencies() as c:

        def true_fn():
          v.assign(v + 1, name="true")
          return 1.0

        def false_fn():

          def inner_true_fn():
            v.assign(v * 2, name="false_true")
            return 2.0

          def inner_false_fn():
            v.assign(v * 3, name="false_false")
            return 3.0

          control_flow_ops.cond(q, inner_true_fn, inner_false_fn)
          return 1.0

        control_flow_ops.cond(p, true_fn, false_fn)
        with ops.name_scope("final"):
          val = v.read_value()
        val = c.mark_as_return(val)
      self.assertAllEqual(val.eval(feed_dict={p: False, q: False}), 3.0)
      self.assertAllEqual(val.eval(feed_dict={p: False, q: True}), 6.0)
      self.assertAllEqual(val.eval(feed_dict={p: True, q: True}), 7.0)
      self.assertAllEqual(val.eval(feed_dict={p: True, q: False}), 8.0)

  @test_util.run_v1_only("b/120545219")
  def testCondOneBranch(self):
    with context.graph_mode(), self.cached_session():
      v = resource_variable_ops.ResourceVariable(1.0)
      self.evaluate(variables.global_variables_initializer())
      p = array_ops.placeholder(dtype=dtypes.bool)
      with acd.AutomaticControlDependencies() as c:

        def true_fn():
          return 0.0

        def false_fn():
          v.assign(v + 4)
          return 1.0

        control_flow_ops.cond(p, true_fn, false_fn)
        val = v.read_value()
        val = c.mark_as_return(val)
      self.assertAllEqual(val.eval(feed_dict={p: False}), 5.0)
      self.assertAllEqual(val.eval(feed_dict={p: True}), 5.0)

  @test_util.run_v1_only("b/120545219")
  def testCondOneBranchUpdateBefore(self):
    with context.graph_mode(), self.cached_session():
      v = resource_variable_ops.ResourceVariable(1.0)
      self.evaluate(variables.global_variables_initializer())
      p = array_ops.placeholder(dtype=dtypes.bool)
      with acd.AutomaticControlDependencies() as c:
        v.assign(v * 2)

        def true_fn():
          return 0.0

        def false_fn():
          v.assign(v + 4)
          return 1.0

        control_flow_ops.cond(p, true_fn, false_fn)
        val = v.read_value()
        val = c.mark_as_return(val)
      self.assertAllEqual(val.eval(feed_dict={p: False}), 6.0)
      self.assertAllEqual(val.eval(feed_dict={p: True}), 12.0)

  @test_util.run_v1_only("b/120545219")
  def testCondOneBranchUpdateAfter(self):
    with context.graph_mode(), self.cached_session():
      v = resource_variable_ops.ResourceVariable(1.0)
      self.evaluate(variables.global_variables_initializer())
      p = array_ops.placeholder(dtype=dtypes.bool)
      with acd.AutomaticControlDependencies() as c:

        def true_fn():
          return 0.0

        def false_fn():
          v.assign(v + 4)
          return 1.0

        control_flow_ops.cond(p, true_fn, false_fn)
        v.assign(v * 2)
        val = v.read_value()
        val = c.mark_as_return(val)
      self.assertAllEqual(val.eval(feed_dict={p: False}), 10.0)
      self.assertAllEqual(val.eval(feed_dict={p: True}), 20.0)

  def testDefunWhileLoopWithCapturedLoopVars(self):
    n = 3
    x = constant_op.constant(list(range(n)))

    @function.defun
    def loop():
      c = lambda i, x: i < n
      b = lambda i, x: (i + 1, x + 1)
      i, out = control_flow_ops.while_loop(c, b, (0, x))
      return i, out

    i, out = loop()
    self.assertEqual(int(i), 3)
    self.assertAllEqual(out, [3, 4, 5])

  def testDecorator(self):
    with context.graph_mode(), self.cached_session():
      v = resource_variable_ops.ResourceVariable(1.0)
      self.evaluate(variables.global_variables_initializer())

      @acd.automatic_control_dependencies
      def f():
        v.assign(v + 1)
        v.assign(2 * v)
        return v.read_value()

      self.assertAllEqual(f(), 4.0)

  def testOptimizerInDefun(self):
    def loss(v):
      return v**2

    optimizer = momentum.MomentumOptimizer(learning_rate=1.0, momentum=1.0)

    @function.defun
    def train():
      self.v = resource_variable_ops.ResourceVariable(1.0)
      grad = backprop.implicit_grad(loss)(self.v)
      optimizer.apply_gradients(grad)
      return self.v.read_value()

    value = train()
    self.assertEqual(value.numpy(), -1.0)

  def testReturningNonTensorRaisesError(self):
    optimizer = momentum.MomentumOptimizer(learning_rate=1.0, momentum=1.0)
    optimizer.apply_gradients = function.defun(optimizer.apply_gradients)
    v = resource_variable_ops.ResourceVariable(1.0)
    grad = backprop.implicit_grad(lambda v: v**2)(v)

    with self.assertRaisesRegex(TypeError,
                                ".*must return zero or more Tensors.*"):
      # TODO(akshayka): We might want to allow defun-ing Python functions
      # that return operations (and just execute the op instead of running it).
      optimizer.apply_gradients(grad)

  # TODO(b/111663004): This should work when the outer context is graph
  # building.
  def testOptimizerNonSlotVarsInDefunNoError(self):
    def loss(v):
      return v**2

    optimizer = adam.AdamOptimizer(learning_rate=1.0)

    @function.defun
    def train():
      self.v = resource_variable_ops.ResourceVariable(1.0)
      grad = backprop.implicit_grad(loss)(self.v)
      optimizer.apply_gradients(grad)
      return self.v.read_value()

    train()

  def testOptimizerInDefunWithCapturedVariable(self):
    v = resource_variable_ops.ResourceVariable(1.0)
    def loss():
      return v**2

    optimizer = momentum.MomentumOptimizer(learning_rate=1.0, momentum=1.0)

    @function.defun
    def train():
      grad = backprop.implicit_grad(loss)()
      optimizer.apply_gradients(grad)

    train()
    self.assertEqual(v.numpy(), -1.0)

  def testRepeatedResourceInput(self):
    var = resource_variable_ops.ResourceVariable(1.0)

    @def_function.function
    def inner(var1, var2):
      return (resource_variable_ops.read_variable_op(var1, dtypes.float32) +
              resource_variable_ops.read_variable_op(var2, dtypes.float32))

    @def_function.function
    def outer():
      return inner(var.handle, var.handle)

    self.assertEqual(self.evaluate(outer()), 2.0)


if __name__ == "__main__":
  ops.enable_eager_execution()
  test.main()
