load_test.py

# 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.
# ==============================================================================
"""Tests for trackable object SavedModel loading."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import os
import sys
import tempfile
import weakref

from absl.testing import parameterized

from tensorflow.python.client import session as session_lib
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.eager import backprop
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
from tensorflow.python.eager import test
from tensorflow.python.eager import wrap_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import function as framework_function
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_spec
from tensorflow.python.framework import test_util
from tensorflow.python.framework import versions
from tensorflow.python.lib.io import file_io
from tensorflow.python.module import module
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import cond_v2
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import lookup_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import numpy_ops as tnp
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables
from tensorflow.python.ops.numpy_ops import np_arrays
from tensorflow.python.ops.ragged import ragged_factory_ops
from tensorflow.python.ops.ragged import ragged_tensor
from tensorflow.python.saved_model import load
from tensorflow.python.saved_model import load_options
from tensorflow.python.saved_model import save
from tensorflow.python.saved_model import tag_constants
from tensorflow.python.training import monitored_session
from tensorflow.python.training.tracking import tracking
from tensorflow.python.training.tracking import util
from tensorflow.python.util import tf_inspect


def cycle(obj, cycles, signatures=None):
  to_save = obj
  # TODO(vbardiovsky): It would be nice if exported protos reached a fixed
  # point w.r.t. saving/restoring, ideally after 2nd saving.
  for _ in range(cycles):
    path = tempfile.mkdtemp(prefix=test.get_temp_dir())
    # If available, we'll run the save and restore preferring the GPU. This
    # just makes sure we aren't throwing errors and have enough
    # device("CPU") blocks to satisfy the placer.
    with test_util.use_gpu():
      save.save(to_save, path, signatures)
      loaded = load.load(path)
    to_save = loaded
  return loaded


@parameterized.named_parameters(
    dict(testcase_name="ReloadOnce", cycles=1),
    dict(testcase_name="ReloadTwice", cycles=2),
    dict(testcase_name="ReloadThrice", cycles=3))
class LoadTest(test.TestCase, parameterized.TestCase):

  def test_structure_import(self, cycles):
    root = tracking.AutoTrackable()
    root.dep_one = tracking.AutoTrackable()
    root.dep_two = tracking.AutoTrackable()
    root.dep_two.dep = tracking.AutoTrackable()
    root.dep_three = root.dep_two.dep
    imported = cycle(root, cycles)
    self.assertIs(imported.dep_three, imported.dep_two.dep)
    self.assertIsNot(imported.dep_one, imported.dep_two)

  @test_util.run_in_graph_and_eager_modes
  def test_variables(self, cycles):
    root = tracking.AutoTrackable()
    root.v1 = variables.Variable(1., trainable=True)
    root.v2 = variables.Variable(2., trainable=False)
    self.evaluate([root.v1.initializer, root.v2.initializer])

    for _ in range(cycles):
      imported = cycle(root, 1)
      self.evaluate([imported.v1.initializer, imported.v2.initializer])

    if not context.executing_eagerly():
      self.assertIsInstance(imported.v1.initializer, ops.Operation)
      self.assertIsInstance(imported.v2.initializer, ops.Operation)

    self.assertEqual(self.evaluate(imported.v1), 1.0)
    self.assertTrue(imported.v1.trainable)
    self.assertEqual(self.evaluate(imported.v2), 2.0)
    self.assertFalse(imported.v2.trainable)

  def test_variables_name(self, cycles):
    root = tracking.AutoTrackable()
    # Test 2 variables with same name: should work as the checkpoint
    # is based on object name and not on variable name.
    root.v1 = variables.Variable(1., trainable=True, name="v1")
    root.v2 = variables.Variable(2., trainable=False, name="v1")
    imported = cycle(root, cycles)
    self.assertEqual(imported.v1.numpy(), 1.0)
    self.assertEqual(imported.v2.numpy(), 2.0)
    self.assertEqual(imported.v1.name, root.v1.name)
    self.assertEqual(imported.v2.name, root.v2.name)
    with variable_scope.variable_scope("foo"):
      imported = cycle(root, cycles)
      self.assertTrue(imported.v1.name.startswith("foo/"))
      self.assertTrue(imported.v2.name.startswith("foo/"))

  def test_partially_defined_variable_shape(self, cycles):

    class MakeVariable(module.Module):

      def __init__(self):
        self.v = None

      @def_function.function(
          input_signature=[tensor_spec.TensorSpec([None], dtypes.int64)])
      def make_variable(self, initial_value):
        if self.v is None:
          self.v = variables.Variable(initial_value)

    m = MakeVariable()
    m.make_variable([1, 2, 3])
    m = cycle(m, cycles)
    m.v.assign([1, 2, 3, 4])
    self.assertEqual([None], tensor_shape.as_shape(m.v.shape).as_list())

  @test_util.run_in_graph_and_eager_modes
  def test_capture_variables(self, cycles):
    root = tracking.AutoTrackable()
    root.weights = variables.Variable(2.)
    self.evaluate(root.weights.initializer)
    root.f = def_function.function(
        lambda x: root.weights * x,
        input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
    for _ in range(cycles):
      imported = cycle(root, 1)
      self.evaluate(imported.weights.initializer)
    self.assertEqual(4., self.evaluate(imported.f(constant_op.constant(2.))))
    self.evaluate(imported.weights.assign(4.0))
    self.assertEqual(8., self.evaluate(imported.f(constant_op.constant(2.))))

  @test_util.run_in_graph_and_eager_modes
  def test_capture_constant(self, cycles):
    root = tracking.AutoTrackable()
    captured_constant = constant_op.constant(2.)
    root.f = def_function.function(
        lambda x: captured_constant * x,
        input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
    imported = cycle(root, cycles)
    self.assertEqual(4., self.evaluate(imported.f(constant_op.constant(2.))))

  def test_control_outputs(self, cycles):
    exported = tracking.AutoTrackable()
    exported.v = variables.Variable(1.)
    exported.f = def_function.function(
        lambda: exported.v.assign(2., name="should_be_control_output"))
    exported_graph = exported.f.get_concrete_function().graph
    self.assertIn(
        exported_graph.get_operation_by_name("should_be_control_output"),
        exported_graph.control_outputs)

    imported = cycle(exported, cycles)
    # Calling get_concrete_function wraps in a second call operation; we want to
    # inspect the original function body for the control output; digging into
    # graph.as_graph_def() and its FunctionDefLibrary is another option.
    imported_concrete, = imported.f.concrete_functions
    imported_graph = imported_concrete.graph
    self.assertIn(
        imported_graph.get_operation_by_name("should_be_control_output"),
        imported_graph.control_outputs)

  def _make_asset(self, contents):
    filename = tempfile.mktemp(prefix=self.get_temp_dir())
    with open(filename, "w") as f:
      f.write(contents)
    return filename

  @test_util.run_in_graph_and_eager_modes
  def test_assets(self, cycles):
    file1 = self._make_asset("contents 1")
    file2 = self._make_asset("contents 2")

    root = tracking.AutoTrackable()
    root.asset1 = tracking.Asset(file1)
    root.asset2 = tracking.Asset(file2)

    save_dir = os.path.join(self.get_temp_dir(), "save_dir")
    save.save(root, save_dir)

    file_io.delete_file(file1)
    file_io.delete_file(file2)
    load_dir = os.path.join(self.get_temp_dir(), "load_dir")
    file_io.rename(save_dir, load_dir)

    imported = load.load(load_dir)
    with open(self.evaluate(imported.asset1.asset_path), "r") as f:
      self.assertEqual("contents 1", f.read())
    with open(self.evaluate(imported.asset2.asset_path), "r") as f:
      self.assertEqual("contents 2", f.read())

  def test_cond_prune(self, cycles):
    x_in = []
    x_out = []

    def f(x, y):
      x_in.append(x)
      xx = cond_v2.cond_v2(
          math_ops.less(1, 2),
          lambda: x + 1,
          lambda: x + 2,
      )
      x_out.append(xx)
      return xx, 2 * y

    f_wrapped = wrap_function.wrap_function(
        f, [tensor_spec.TensorSpec((), dtypes.float32)] * 2)
    f_pruned = f_wrapped.prune(x_in[0], [x_out[0]])

    class Adder(module.Module):

      @def_function.function(input_signature=[
          tensor_spec.TensorSpec(shape=None, dtype=dtypes.float32)])
      def add(self, x):
        return f_pruned(x)

    root = Adder()
    root.add(constant_op.constant(1.))
    root = cycle(root, cycles)
    root.add(constant_op.constant(1.))

  def test_capture_assets(self, cycles):
    root = tracking.AutoTrackable()
    root.vocab = tracking.Asset(self._make_asset("contents"))
    root.f = def_function.function(
        lambda: root.vocab.asset_path,
        input_signature=[])
    imported = cycle(root, cycles)
    original_output = root.f().numpy()
    imported_output = imported.f().numpy()
    self.assertNotEqual(original_output, imported_output)
    with open(imported_output, "r") as f:
      self.assertEqual("contents", f.read())

  def test_capture_assets_in_graph(self, cycles):
    root = tracking.AutoTrackable()
    root.vocab = tracking.Asset(self._make_asset("contents"))
    root.f = def_function.function(
        lambda: root.vocab.asset_path,
        input_signature=[])

    original_output = root.f().numpy()

    if cycles > 1:
      root = cycle(root, cycles - 1)
    path = tempfile.mkdtemp(prefix=self.get_temp_dir())
    save.save(root, path)

    with ops.Graph().as_default():
      imported = load.load(path)
      imported_tensor = imported.f()
      with monitored_session.MonitoredSession() as sess:
        imported_output = sess.run(imported_tensor)
        self.assertNotEqual(original_output, imported_output)
        with open(imported_output, "r") as f:
          self.assertEqual("contents", f.read())

  def test_dedup_assets(self, cycles):
    vocab = self._make_asset("contents")
    root = tracking.AutoTrackable()
    root.asset1 = tracking.Asset(vocab)
    root.asset2 = tracking.Asset(vocab)
    imported = cycle(root, cycles)
    self.assertEqual(imported.asset1.asset_path.numpy(),
                     imported.asset2.asset_path.numpy())

  def test_implicit_input_signature(self, cycles):
    @def_function.function
    def func(x):
      return 2 * x

    root = tracking.AutoTrackable()
    root.f = func

    # Add two traces.
    root.f(constant_op.constant(1.))
    root.f(constant_op.constant(1))

    imported = cycle(root, cycles)

    self.assertEqual(4., imported.f(constant_op.constant(2.)).numpy())
    self.assertEqual(14, imported.f(constant_op.constant(7)).numpy())

  def test_explicit_input_signature(self, cycles):
    @def_function.function(
        input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
    def func(x):
      return 2 * x

    root = tracking.AutoTrackable()
    root.f = func

    imported = cycle(root, cycles)
    self.assertEqual(4., imported.f(constant_op.constant(2.0)).numpy())

  def test_explicit_save_signature(self, cycles):
    @def_function.function
    def func(x):
      return 2 * x

    root = tracking.AutoTrackable()
    root.f = func

    imported = cycle(
        root, cycles, {
            "f":
                root.f.get_concrete_function(
                    tensor_spec.TensorSpec(None, dtypes.float32))
        })
    self.assertEqual(4., imported.f(constant_op.constant(2.0)).numpy())

  def test_nested_functions(self, cycles):
    f = def_function.function(
        lambda x: x*2.0,
        input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
    g = def_function.function(
        lambda x: f(x) + 1.0,
        input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])

    root = tracking.AutoTrackable()
    root.g = g
    imported = cycle(root, cycles)
    imported.g(constant_op.constant([1.0]))

  def test_function_with_default_bool_input(self, cycles):

    def func(x, training=False):
      if training:
        return 2 * x
      else:
        return 7

    root = tracking.AutoTrackable()
    root.f = def_function.function(func)

    self.assertEqual(20, root.f(constant_op.constant(10), True).numpy())
    self.assertEqual(7, root.f(constant_op.constant(1)).numpy())
    self.assertEqual(2, root.f(constant_op.constant(1), True).numpy())

    imported = cycle(root, cycles)

    self.assertEqual(4, imported.f(constant_op.constant(2), True).numpy())
    self.assertEqual(7, imported.f(constant_op.constant(2)).numpy())

  def test_function_with_default_none_input(self, cycles):

    def func(x, dtype=None):
      if dtype:
        return array_ops.zeros(shape=x.shape, dtype=dtype)
      else:
        return array_ops.zeros(shape=x.shape, dtype=dtypes.float32)

    root = tracking.AutoTrackable()
    root.f = def_function.function(func)

    self.assertAllEqual([0.0, 0.0, 0.0],
                        root.f(constant_op.constant([1, 2, 3])).numpy())
    self.assertAllEqual([0.0, 0.0, 0.0],
                        root.f(constant_op.constant([1.0, 2.0, 3.0])).numpy())
    self.assertAllEqual([0.0, 0.0, 0.0, 0.0],
                        root.f(constant_op.constant([1, 2, 3, 4])).numpy())
    self.assertAllEqual([0, 0, 0],
                        root.f(
                            constant_op.constant([1.0, 2.0, 3.0]),
                            dtype=dtypes.int32).numpy())

    concrete_functions = root.f._list_all_concrete_functions_for_serialization()  # pylint: disable=protected-access
    self.assertEqual(4, len(concrete_functions))

    imported = cycle(root, cycles)

    self.assertAllEqual([0.0, 0.0, 0.0],
                        imported.f(constant_op.constant([1, 2, 3]),
                                   None).numpy())
    self.assertAllEqual([0.0, 0.0, 0.0],
                        imported.f(constant_op.constant([1.0, 2.0,
                                                         3.0])).numpy())
    self.assertAllEqual([0.0, 0.0, 0.0, 0.0],
                        imported.f(constant_op.constant([1, 2, 3, 4])).numpy())
    self.assertAllEqual([0, 0, 0],
                        imported.f(
                            constant_op.constant([1.0, 2.0, 3.0]),
                            dtype=dtypes.int32).numpy())

  def test_function_no_return(self, cycles):

    class TrackableWithOneVariable(tracking.AutoTrackable):

      def __init__(self, initial_value=0.0):
        super(TrackableWithOneVariable, self).__init__()
        self.variable = variables.Variable(initial_value)

      @def_function.function
      def increase(self, by=1.0):
        self.variable.assign_add(by)

    obj = TrackableWithOneVariable(5.0)

    obj.increase(constant_op.constant(10.0))
    self.assertEqual(15.0, obj.variable.numpy())
    obj.increase()
    self.assertEqual(16.0, obj.variable.numpy())

    imported = cycle(obj, cycles)

    imported.increase(constant_op.constant(10.0))
    self.assertEqual(26.0, imported.variable.numpy())
    imported.increase(constant_op.constant(1.0))
    self.assertEqual(27.0, imported.variable.numpy())

  def test_structured_inputs(self, cycles):

    def func(x, training=True):
      # x is a nested structure, we care about one particular tensor.
      _, (a, b) = x
      if training:
        return 2 * a["a"] + b
      else:
        return 7

    root = tracking.AutoTrackable()
    root.f = def_function.function(func)

    x = constant_op.constant(10)
    y = constant_op.constant(11)

    input1 = [6, ({"a": x}, y)]
    input2 = [7, ({"a": x}, y)]  # Not compatible with input1 signature.
    input3 = [6, ({"a": y}, x)]  # Compatible with input1 signature.

    # Note: by only calling f(input1) before serialization, only inputs with
    # matching signature will be valid on the loaded model.
    self.assertEqual(31, root.f(input1).numpy())

    imported = cycle(root, cycles)

    with self.assertRaisesRegex(ValueError,
                                "Could not find matching function to call"):
      imported.f(input2)

    self.assertEqual(31, imported.f(input1).numpy())
    self.assertEqual(32, imported.f(input3).numpy())

  def test_structured_output(self, cycles):

    # Use fields with non-alphabetical order
    named_tuple_type = collections.namedtuple("NamedTupleHello", ["b", "a"])

    def func(input1, input2):
      named_tuple = named_tuple_type(a=input1 + input2, b=input1 * input2)
      return [named_tuple, input2, {"x": 0.5}]

    root = tracking.AutoTrackable()
    root.f = def_function.function(func)

    result = root.f(constant_op.constant(2), constant_op.constant(3))

    self.assertEqual(5, result[0].a.numpy())
    self.assertEqual(6, result[0].b.numpy())
    self.assertEqual(["b", "a"], list(result[0]._asdict().keys()))
    self.assertEqual(3, result[1].numpy())
    self.assertEqual(0.5, result[2]["x"].numpy())

    imported = cycle(root, cycles)

    result = imported.f(constant_op.constant(2), constant_op.constant(5))
    self.assertEqual(7, result[0].a.numpy())
    self.assertEqual(10, result[0].b.numpy())
    self.assertEqual(["b", "a"], list(result[0]._asdict().keys()))
    self.assertEqual(5, result[1].numpy())
    self.assertEqual(0.5, result[2]["x"].numpy())

  def test_positional_arguments(self, cycles):
    def func(x, training=False, abc=7.1, defg=7.7):
      del abc
      if training:
        return 2 * x
      if defg == 7:
        return 6
      else:
        return 7

    root = tracking.AutoTrackable()
    root.f = def_function.function(func)

    self.assertEqual(20, root.f(constant_op.constant(10), True).numpy())
    self.assertEqual(7, root.f(constant_op.constant(1)).numpy())
    self.assertEqual(2, root.f(constant_op.constant(1), True).numpy())
    self.assertEqual(6, root.f(constant_op.constant(1), defg=7.0).numpy())

    imported = cycle(root, cycles)

    self.assertEqual(4, imported.f(constant_op.constant(2), True).numpy())
    self.assertEqual(7, imported.f(constant_op.constant(2)).numpy())
    self.assertEqual(6, imported.f(constant_op.constant(1), defg=7.0).numpy())

  def test_additional_kwargs(self, cycles):
    def func(x, training=False, **options):
      del options
      if training:
        return 2 * x
      else:
        return 7

    root = tracking.AutoTrackable()
    root.f = def_function.function(func)

    x = constant_op.constant(10)
    self.assertEqual(7, root.f(x, learning_rate=0.5, epochs=3).numpy())

    imported = cycle(root, cycles)

    with self.assertRaisesRegex(ValueError,
                                "Could not find matching function to call.*"):
      imported.f(x, learning_rate=0.5, epochs=4)

    self.assertEqual(7, imported.f(x, learning_rate=0.5, epochs=3).numpy())

  def test_member_function(self, cycles):
    class TrackableWithMember(tracking.AutoTrackable):

      def __init__(self):
        super(TrackableWithMember, self).__init__()
        self._some_value = 20

      @def_function.function
      def f(self, x, training=False):
        if training:
          return 2 * x
        else:
          return 7 + self._some_value

    root = TrackableWithMember()

    self.assertEqual(20, root.f(constant_op.constant(10), True).numpy())
    self.assertEqual(27, root.f(constant_op.constant(1)).numpy())
    self.assertEqual(2, root.f(constant_op.constant(1), True).numpy())

    imported = cycle(root, cycles)

    self.assertEqual(4, imported.f(constant_op.constant(2), True).numpy())
    self.assertEqual(27, imported.f(constant_op.constant(2)).numpy())

  def test_side_effect_listing(self, cycles):
    class M(tracking.AutoTrackable):

      def __init__(self):
        super(M, self).__init__()
        self.var = None

      @def_function.function(
          input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
      def f(self, x):
        if self.var is None:
          self.var = variables.Variable(2.)
        return x * self.var

    m = M()
    cycle(m, cycles)
    self.assertEqual(4.0, m.f(constant_op.constant(2.0)).numpy())

  def test_basic_backprop(self, cycles):
    weight = variables.Variable(1., trainable=True)
    bias = variables.Variable(0., trainable=True)
    g = def_function.function(
        lambda x: x*weight + bias,
        input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])

    root = tracking.AutoTrackable()
    root.weight = weight
    root.bias = bias
    root.g = g
    imported = cycle(root, cycles)
    with backprop.GradientTape() as t:
      x = constant_op.constant([3.5])
      loss = imported.g(x)
      grad = t.gradient(loss, [imported.weight, imported.bias])
      self.assertAllClose(grad, [3.5, 1.0])

  def test_nested_backprop(self, cycles):
    weight = variables.Variable(1., trainable=True)
    bias = variables.Variable(0., trainable=True)

    # Note: this function gets called from other function defs via a
    # "PartitionedCall" op node.
    @def_function.function(input_signature=[
        tensor_spec.TensorSpec(None, dtypes.float32),
        tensor_spec.TensorSpec(None, dtypes.float32)])
    def mul(x, y):
      return x * y

    # Note: this function gets called from other function defs via a
    # "StatefulPartitionedCall" op node.
    @def_function.function(input_signature=[
        tensor_spec.TensorSpec(None, dtypes.float32)])
    def f(x):
      return mul(weight.read_value(), x)

    @def_function.function(input_signature=[
        tensor_spec.TensorSpec(None, dtypes.float32)])
    def g(x):
      return f(x) + bias,

    @def_function.function(input_signature=[
        tensor_spec.TensorSpec(None, dtypes.float32)])
    def h(x):
      return g(x) + bias,

    root = tracking.AutoTrackable()
    root.weight = weight
    root.bias = bias
    root.g = h

    imported = cycle(root, cycles)
    with backprop.GradientTape() as t:
      x = constant_op.constant([3.5])
      loss = imported.g(x)
    grad = t.gradient(loss, [imported.weight, imported.bias])
    self.assertAllClose(grad, [3.5, 2.0])

  def test_while_loop_backprop(self, cycles):
    weight = variables.Variable(2., trainable=True)

    @def_function.function(input_signature=[
        tensor_spec.TensorSpec(dtype=dtypes.float32, shape=(None, None))])
    def g(x):
      """Adds rows of matrix x after multiplying each entry by v."""
      i_0 = constant_op.constant(0)
      s_0 = constant_op.constant([0., 0.])
      cond = lambda i, _: i < array_ops.shape(x)[1]
      body = lambda i, s: (i + 1, s + weight * x[:, i])
      i_end, s_end = control_flow_ops.while_loop(cond, body, (i_0, s_0))
      del i_end
      return s_end

    root = tracking.AutoTrackable()
    root.weight = weight
    root.g = g
    imported = cycle(root, cycles)

    def get_gradient(obj):
      with backprop.GradientTape() as t:
        x = constant_op.constant([[1., 2., 3.], [1., -2, 3.]])
        y = obj.g(x)
        self.assertAllClose(y, obj.weight * [6., 2.])
        loss = math_ops.reduce_sum(y)  # weight * 8.
        self.assertAllEqual(t.watched_variables(), [obj.weight])
        return t.gradient(loss, obj.weight)

    imported_gradient = get_gradient(imported)
    original_gradient = get_gradient(root)
    self.assertIsNotNone(original_gradient)
    self.assertAllClose(original_gradient, 8.)
    self.assertIsNotNone(imported_gradient)
    self.assertAllClose(imported_gradient, 8.)

  def _test_restored_func_with_captured_var_backprop(self, cycles, dtype):
    weight = variables.Variable(2., trainable=True, dtype=dtype)

    @def_function.function(input_signature=[
        tensor_spec.TensorSpec(dtype=dtype, shape=())])
    def g(x):
      return x * weight

    root = tracking.AutoTrackable()
    root.weight = weight
    root.g = g
    imported = cycle(root, cycles)

    def get_gradient(obj):
      with backprop.GradientTape() as t:
        x = constant_op.constant(2.)
        y = obj.g(x)
        self.assertAllClose(y, obj.weight * 2.)
        self.assertAllEqual(t.watched_variables(), [obj.weight])
        return t.gradient(y, obj.weight)

    imported_gradient = get_gradient(imported)
    original_gradient = get_gradient(root)
    self.assertIsNotNone(original_gradient)
    self.assertAllClose(original_gradient, 2.)
    self.assertIsNotNone(imported_gradient)
    self.assertAllClose(imported_gradient, 2.)

  def test_restored_func_with_captured_var_backprop_float32(self, cycles):
    self._test_restored_func_with_captured_var_backprop(cycles, dtypes.float32)

  def test_restored_func_with_captured_var_backprop_float64(self, cycles):
    self.skipTest("b/144573917")
    self._test_restored_func_with_captured_var_backprop(cycles, dtypes.float64)

  def test_callable(self, cycles):
    class M1(tracking.AutoTrackable):

      @def_function.function(
          input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
      def __call__(self, x):
        return x

    root = tracking.AutoTrackable()
    root.m1 = M1()
    root.m2 = tracking.AutoTrackable()
    root.m2.__call__ = def_function.function(
        input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])(
            lambda x: x*3.0)
    imported = cycle(root, cycles)
    x = constant_op.constant(1.0)

    self.assertTrue(callable(imported.m1))
    self.assertAllEqual(root.m1(x), imported.m1(x))

    # Note: `root.m2` was not callable since `__call__` attribute was set
    # into the instance and not on the class. But after a serialization cycle
    # that starts to work.
    self.assertTrue(callable(imported.m2))
    self.assertAllEqual(root.m2.__call__(x), imported.m2(x))

    # Verify that user objects without `__call__` attribute are not callable.
    self.assertFalse(callable(imported))

  def test_chain_callable(self, cycles):
    func = def_function.function(
        input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])(
            lambda x: x*3.0)
    root = tracking.AutoTrackable()
    root.__call__ = tracking.AutoTrackable()
    root.__call__.__call__ = tracking.AutoTrackable()
    root.__call__.__call__.__call__ = func

    imported = cycle(root, cycles)
    self.assertTrue(callable(imported))
    x = constant_op.constant(1.0)
    self.assertAllEqual(imported(x).numpy(), 3.0)

  def test_load_in_graph_mode(self, cycles):
    root = tracking.AutoTrackable()
    root.v1 = variables.Variable(1., name="v_one", trainable=False)
    root.v2 = variables.Variable(2., name="v_two", trainable=True)
    root.f = def_function.function(
        lambda x: root.v2 * x,
        input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])

    if cycles > 1:
      root = cycle(root, cycles - 1)
    path = tempfile.mkdtemp(prefix=self.get_temp_dir())
    save.save(root, path)

    with ops.Graph().as_default() as g:
      imported = load.load(path)
      var_v1 = imported.v1
      self.assertFalse(var_v1.trainable)
      var_v2 = imported.v2
      self.assertTrue(var_v2.trainable)
      output = imported.f(constant_op.constant(2.))
      with monitored_session.MonitoredSession() as sess:
        self.assertEqual(1.0, sess.run(var_v1))
        self.assertEqual(4.0, sess.run(output))
      self.assertCountEqual([var_v1, var_v2],
                            g.get_collection(ops.GraphKeys.GLOBAL_VARIABLES))
      # load() should not add to TRAINABLE_VARIABLES. Higher levels of model
      # building control retraining or frozen use of imported SavedModels.
      self.assertCountEqual([],
                            g.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES))

  def test_load_in_func_graph(self, cycles):
    root = tracking.AutoTrackable()
    root.v1 = variables.Variable(1.)
    root.v2 = variables.Variable(2.)
    root.f = def_function.function(
        lambda x: root.v2 * x,
        input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])

    if cycles > 1:
      root = cycle(root, cycles - 1)
    path = tempfile.mkdtemp(prefix=self.get_temp_dir())
    save.save(root, path)

    closure = tracking.AutoTrackable()
    @def_function.function
    def func(x):
      if not hasattr(closure, "model"):
        closure.model = load.load(path)
      return closure.model.f(x)

    inputs = constant_op.constant(2.)
    self.assertEqual(4.0, func(inputs).numpy())

  def test_soft_matching(self, cycles):

    @def_function.function(
        input_signature=[tensor_spec.TensorSpec([None], dtypes.int32)])
    def func(x):
      return 2 * x

    root = tracking.AutoTrackable()
    root.f = func

    self.assertAllEqual([2], root.f(constant_op.constant([1])).numpy())
    self.assertAllEqual([2, 4], root.f(constant_op.constant([1, 2])).numpy())

    concrete_functions = root.f._list_all_concrete_functions_for_serialization()  # pylint: disable=protected-access
    self.assertEqual(1, len(concrete_functions))

    imported = cycle(root, cycles)

    with self.assertRaisesRegex(ValueError, "Python inputs incompatible"):
      # We cannot call the function with a constant of shape ().
      imported.f(constant_op.constant(2)).numpy()

    # TODO(vbardiovsky): When classes are revived with input_signatures, we
    # should also check that the calls below are not generating any more
    # concrete functions.
    self.assertAllEqual([2, 4, 6, 8],
                        imported.f(constant_op.constant([1, 2, 3, 4])).numpy())
    self.assertAllEqual([2, 4, 6],
                        imported.f(constant_op.constant([1, 2, 3])).numpy())

  def test_get_concrete_function(self, cycles):

    @def_function.function
    def func(x, training=False):
      if training:
        return 2 * x
      else:
        return 3 * x

    func.get_concrete_function(
        tensor_spec.TensorSpec([None], dtypes.int32), True)
    func.get_concrete_function(tensor_spec.TensorSpec([None], dtypes.float32))

    root = tracking.AutoTrackable()
    root.f = func

    imported = cycle(root, cycles)

    concrete = imported.f.get_concrete_function(
        training=True, x=tensor_spec.TensorSpec([None], dtypes.int32))

    self.assertAllEqual([2, 4, 6, 8],
                        concrete(x=constant_op.constant([1, 2, 3, 4])).numpy())
    with self.assertRaisesRegex(ValueError,
                                "Could not find matching function to call"):
      imported.f.get_concrete_function(
          tensor_spec.TensorSpec([None], dtypes.int32))
    imported.f.get_concrete_function(
        tensor_spec.TensorSpec([None], dtypes.int32), True)

  def test_concrete_function(self, cycles):

    @def_function.function(
        input_signature=[tensor_spec.TensorSpec([None], dtypes.int32)])
    def func(x):
      return 2 * x

    root = tracking.AutoTrackable()
    root.f = func.get_concrete_function()

    self.assertAllEqual([2], root.f(constant_op.constant([1])).numpy())
    self.assertAllEqual([2, 4], root.f(constant_op.constant([1, 2])).numpy())

    # TODO(andresp): Fix exporting of loaded concrete functions as signatures.
    imported = cycle(root, cycles, signatures={})

    self.assertAllEqual([2, 4, 6, 8],
                        imported.f(constant_op.constant([1, 2, 3, 4])).numpy())
    self.assertAllEqual([2, 4, 6],
                        imported.f(constant_op.constant([1, 2, 3])).numpy())

  def test_concrete_function_captures(self, cycles):

    class Root(module.Module):

      def __init__(self):
        self.v = variables.Variable(1.)
        self.v1 = variables.Variable(1.)

      @def_function.function(
          input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
      def use_v(self, x):
        return self.v + self.v1 + 1.

    root = Root()
    self.assertIn(root.v.handle,
                  root.use_v.get_concrete_function().graph.external_captures)
    for _ in range(cycles):
      root = cycle(root, 1, signatures=root.use_v.get_concrete_function())
    func_captures = root.use_v.get_concrete_function().graph.external_captures
    self.assertLen(func_captures, 2)
    self.assertTrue(any(root.v.handle is t for t in func_captures))
    self.assertTrue(any(root.v1.handle is t for t in func_captures))
    signature_captures = root.signatures[
        "serving_default"].graph.external_captures
    self.assertLen(signature_captures, 2)
    self.assertTrue(any(root.v.handle is t for t in signature_captures))
    self.assertTrue(any(root.v1.handle is t for t in signature_captures))

  def test_concrete_function_arg_names(self, cycles):

    @def_function.function(
        input_signature=[tensor_spec.TensorSpec([None], dtypes.int32)])
    def func(x):
      return 2 * x

    root = tracking.AutoTrackable()
    root.f = func.get_concrete_function()

    self.assertAllEqual([2], root.f(constant_op.constant([1])).numpy())

    # TODO(andresp): Fix exporting of loaded concrete functions as signatures.
    imported = cycle(root, cycles, signatures={})

    self.assertAllEqual([2, 4, 6],
                        imported.f(x=constant_op.constant([1, 2, 3])).numpy())

  def test_concrete_function_no_signature(self, cycles):
    @def_function.function
    def func(x):
      return 2 * x

    root = tracking.AutoTrackable()
    root.f = func.get_concrete_function(constant_op.constant([1]))
    self.assertAllEqual([4], root.f(constant_op.constant([2])).numpy())
    # TODO(andresp): Fix exporting of loaded concrete functions as signatures.
    imported = cycle(root, cycles, signatures={})
    self.assertAllEqual([6],
                        imported.f(constant_op.constant([3])).numpy())

  @test_util.run_in_graph_and_eager_modes
  def test_concrete_function_backprop(self, cycles):
    @def_function.function(
        input_signature=[tensor_spec.TensorSpec([], dtypes.float32)])
    def func(x):
      return x ** 2.
    root = tracking.AutoTrackable()
    root.f = func.get_concrete_function()

    def _compute_gradient(function):
      with backprop.GradientTape() as tape:
        inp = constant_op.constant(1.)
        tape.watch(inp)
        output = function(inp)
      return tape.gradient(output, inp)

    self.assertAllEqual(2., _compute_gradient(root.f))
    # TODO(andresp): Fix exporting of loaded concrete functions as signatures.
    imported = cycle(root, cycles, signatures={})
    self.assertAllEqual(2., _compute_gradient(imported.f))

  def test_revived_concrete_function_kwargs(self, cycles):

    @def_function.function
    def func(x, y):
      return x * (y + 1.)
    root = tracking.AutoTrackable()
    root.f = func.get_concrete_function(
        tensor_spec.TensorSpec([], dtypes.float32),
        tensor_spec.TensorSpec([], dtypes.float32))
    self.assertEqual(8., root.f(y=constant_op.constant(3.),
                                x=constant_op.constant(2.)).numpy())
    # TODO(andresp): Fix exporting of loaded concrete functions as signatures.
    imported = cycle(root, cycles, signatures={})
    self.assertEqual(8., imported.f(y=constant_op.constant(3.),
                                    x=constant_op.constant(2.)).numpy())

  def test_revived_concrete_function_tensorspec_kwargs(self, cycles):

    @def_function.function
    def func(*args):
      x, y = args
      return x * (y + 1.)
    root = tracking.AutoTrackable()
    root.f = func.get_concrete_function(
        tensor_spec.TensorSpec([], dtypes.float32, name="x"),
        tensor_spec.TensorSpec([], dtypes.float32, name="y"))
    self.assertEqual(8., root.f(y=constant_op.constant(3.),
                                x=constant_op.constant(2.)).numpy())
    imported = cycle(root, cycles, signatures={})
    self.assertEqual(8., imported.f(y=constant_op.constant(3.),
                                    x=constant_op.constant(2.)).numpy())

  def test_concrete_function_variable_argument(self, cycles):
    capture = variables.Variable(0)

    @def_function.function
    def func(v):
      v.assign_add(1)
      capture.assign_sub(1)

    @def_function.function(input_signature=[
        resource_variable_ops.VariableSpec(shape=[], dtype=dtypes.int32)
    ])
    def func_with_input_signature(v):
      v.assign_add(5)
      capture.assign_sub(5)
      return 1

    vsave = variables.Variable(1)
    root = tracking.AutoTrackable()
    root.f = func.get_concrete_function(vsave)
    root.f_sig = func_with_input_signature.get_concrete_function()
    root.capture = capture

    self.assertEqual(1, vsave.numpy())
    root.f(vsave)
    self.assertEqual(2, vsave.numpy())
    self.assertEqual(-1, capture.numpy())

    root.f_sig(vsave)
    self.assertEqual(7, vsave.numpy())
    self.assertEqual(-6, capture.numpy())

    imported = cycle(root, cycles)

    vload = variables.Variable(1)
    imported.f(vload)
    self.assertEqual(2, vload.numpy())
    imported.f(v=vload)
    self.assertEqual(3, vload.numpy())
    self.assertEqual(-8, imported.capture.numpy())

    imported.f_sig(v=vload)
    self.assertEqual(8, vload.numpy())
    self.assertEqual(-13, imported.capture.numpy())

    self.assertEqual(-6, capture.numpy())

  def test_function_and_component(self, cycles):

    @def_function.function
    def func(v):
      return v + 1

    root = tracking.AutoTrackable()
    root.func = func
    root.concrete_func = func.get_concrete_function(
        tensor_spec.TensorSpec(None, dtypes.int32))
    one = constant_op.constant(1)
    self.assertEqual(2, root.func(one).numpy())
    self.assertEqual(2, root.concrete_func(one).numpy())
    imported = cycle(root, cycles)
    self.assertEqual(2, imported.func(one).numpy())
    self.assertEqual(2, imported.concrete_func(one).numpy())

  def test_dict(self, cycles):
    root = tracking.AutoTrackable()
    root.variables = dict(a=variables.Variable(1.))
    root.variables["b"] = variables.Variable(2.)
    root.variables["c"] = 1
    root.funcs = dict(
        a=def_function.function(lambda: constant_op.constant(100.)))
    root.funcs["conc"] = root.funcs["a"].get_concrete_function()
    imported = cycle(root, cycles)
    self.assertEqual(1., imported.variables["a"].numpy())
    self.assertEqual(2., imported.variables["b"].numpy())
    self.assertEqual(set(["a", "b"]), set(imported.variables.keys()))
    self.assertEqual(100., imported.funcs["a"]().numpy())
    self.assertEqual(100., imported.funcs["conc"]().numpy())

  def test_list(self, cycles):
    root = tracking.AutoTrackable()
    root.variables = [variables.Variable(1.)]
    root.variables.append(1)
    root.variables.append(variables.Variable(3.))
    imported = cycle(root, cycles)
    self.assertEqual(1., imported.variables[0].numpy())
    self.assertEqual(3., imported.variables[2].numpy())
    self.assertIs(None, imported.variables[1])
    self.assertEqual(3, len(imported.variables))

  def test_tuple(self, cycles):
    root = tracking.AutoTrackable()
    root.variables = (variables.Variable(1.), 1, variables.Variable(3.))
    imported = cycle(root, cycles)
    self.assertEqual(1., imported.variables[0].numpy())
    self.assertEqual(3., imported.variables[2].numpy())
    self.assertIs(None, imported.variables[1])
    self.assertLen(imported.variables, 3)

  def test_functions_list(self, cycles):
    root = tracking.AutoTrackable()
    v1 = variables.Variable(1.)
    root.losses = [def_function.function(lambda: math_ops.reduce_sum(v1 ** 2))]
    root.variables = [v1]

    @def_function.function
    def _v2_loss():
      if len(root.variables) == 1:
        v2 = variables.Variable(2.)
        root.variables.append(v2)
      return math_ops.reduce_sum(root.variables[1] ** 2)

    root.losses.append(_v2_loss)
    self.assertAllClose([1., 4.], [loss() for loss in root.losses])
    imported = cycle(root, cycles)
    self.assertAllClose([1., 4.], [loss() for loss in imported.losses])
    imported.variables[0].assign(3.)
    imported.variables[1].assign(4.)
    self.assertAllClose([9., 16.], [loss() for loss in imported.losses])

  def test_captured_constant(self, cycles):
    const = array_ops.zeros([100])
    root = tracking.AutoTrackable()
    root.f = def_function.function(lambda: const + 1.)
    root.g = def_function.function(lambda: const + 2.)
    self.assertAllClose(array_ops.ones([100]), root.f())
    self.assertAllClose(2. * array_ops.ones([100]), root.g())
    imported = cycle(root, cycles)
    self.assertAllClose(array_ops.ones([100]), imported.f())
    self.assertAllClose(2. * array_ops.ones([100]), imported.g())
    # TODO(b/123408994): Use the public get_concrete_function.
    f_concrete = imported.f._list_all_concrete_functions_for_serialization()[0]
    g_concrete = imported.g._list_all_concrete_functions_for_serialization()[0]
    self.assertLen(f_concrete.captured_inputs, 1)
    self.assertLen(g_concrete.captured_inputs, 1)
    # We should be using the same captured EagerTensor in both functions, not
    # duplicating the constant.
    self.assertIs(f_concrete.captured_inputs[0],
                  g_concrete.captured_inputs[0])

  def test_functions_accessed_once(self, cycles):

    class Exported(tracking.AutoTrackable):

      def __init__(self):
        self._counter = 0

      @property
      def make_func(self):
        @def_function.function
        def f():
          return constant_op.constant(self._counter)
        f.get_concrete_function()  # force a trace
        self._counter += 1
        return f

    exported = Exported()
    imported = cycle(exported, cycles)
    self.assertEqual(0, imported.make_func().numpy())
    self.assertEqual(1, exported.make_func().numpy())

  def test_overwritten_signatures_error(self, cycles):
    exported = tracking.AutoTrackable()
    exported.f = def_function.function(lambda: constant_op.constant(1.))
    imported = cycle(
        exported, cycles,
        signatures={"key": exported.f.get_concrete_function()})
    self.assertEqual(1., imported.signatures["key"]()["output_0"].numpy())
    imported.signatures = {"key1": imported.signatures["key"]}
    with self.assertRaisesRegex(ValueError, "signatures"):
      save.save(imported, tempfile.mkdtemp(prefix=self.get_temp_dir()))

  def test_signature_loading(self, cycles):

    class Exported(tracking.AutoTrackable):

      def __init__(self):
        self.v = variables.Variable(3.)

      @def_function.function
      def do(self, x):
        return self.v * x

    exported = Exported()
    imported = cycle(
        exported,
        cycles=1,
        signatures=exported.do.get_concrete_function(
            tensor_spec.TensorSpec(None, dtypes.float32)))
    for _ in range(cycles - 1):
      imported = cycle(imported, cycles=1, signatures=imported.signatures)
    self.assertEqual(["serving_default"], list(imported.signatures.keys()))
    imported_function = imported.signatures["serving_default"]
    two = constant_op.constant(2.)
    self.assertEqual(6., imported_function(x=two)["output_0"].numpy())
    imported.v.assign(4.)
    self.assertEqual(8., imported_function(x=two)["output_0"].numpy())
    self.assertEqual(8., imported_function(two)["output_0"].numpy())
    with self.assertRaises(TypeError):
      # The signatures mapping is immutable
      imported.signatures["random_key"] = 3

  def test_multiple_argument_signatures_no_positional(self, cycles):

    class Exported(tracking.AutoTrackable):

      @def_function.function
      def do(self, x, y):
        return x + y

    exported = Exported()
    imported = cycle(
        exported, cycles=1, signatures=exported.do.get_concrete_function(
            tensor_spec.TensorSpec(None, dtypes.float32),
            tensor_spec.TensorSpec(None, dtypes.float32)))
    for _ in range(cycles - 1):
      imported = cycle(imported, cycles=1, signatures=imported.signatures)
    with self.assertRaises(TypeError):
      imported.signatures["serving_default"](
          constant_op.constant(1.),
          y=constant_op.constant(2.))
    self.assertEqual(
        {"output_0": 3.},
        self.evaluate(imported.signatures["serving_default"](
            x=constant_op.constant(1.),
            y=constant_op.constant(2.))))

  def _make_model_with_tables(self):
    default_val = -1
    keys = constant_op.constant(["brain", "salad", "surgery"])
    values = constant_op.constant([0, 1, 2], dtypes.int64)
    table1_initializer = lookup_ops.KeyValueTensorInitializer(keys, values)
    table1 = lookup_ops.HashTable(table1_initializer, default_val)

    table2_file = self._make_asset("test\nfoo\nbrain\n")
    table2_initializer = lookup_ops.TextFileIdTableInitializer(table2_file)
    table2 = lookup_ops.HashTable(table2_initializer, default_val)

    def _make_lookup_function(table):
      signature = [tensor_spec.TensorSpec(None, dtypes.string)]
      return def_function.function(input_signature=signature)(
          lambda x: table.lookup(x))  # pylint: disable=unnecessary-lambda

    root = tracking.AutoTrackable()
    root.table1 = table1
    root.lookup1 = _make_lookup_function(table1)
    root.table2 = table2
    root.lookup2 = _make_lookup_function(table2)
    return root

  def test_table(self, cycles):
    root = self._make_model_with_tables()
    imported = cycle(root, cycles, signatures={})
    keys = constant_op.constant(["brain", "test", "foo", "surgery"])
    self.assertAllEqual([0, -1, -1, 2], imported.lookup1(keys).numpy())
    self.assertAllEqual([2, 0, 1, -1], imported.lookup2(keys).numpy())

  def test_table_collections_untouched_eager(self, cycles):

    def _gather_nonempty_collections():
      graph = ops.get_default_graph()
      gathered = {}
      for collection in graph.collections:
        collection_contents = graph.get_collection(collection)
        if collection_contents:
          gathered[collection] = collection_contents
      return gathered

    root = self._make_model_with_tables()
    # Warm up collections to ignore those that don't expand every iteration,
    # e.g. the __varscope collection.
    cycle(root, 1)
    original_collections = _gather_nonempty_collections()
    cycle(root, cycles)
    self.assertEqual(original_collections, _gather_nonempty_collections())

  def test_table_in_graph(self, cycles):
    root = self._make_model_with_tables()

    if cycles > 1:
      root = cycle(root, cycles - 1)
    path = tempfile.mkdtemp(prefix=self.get_temp_dir())
    save.save(root, path)
    imported = cycle(root, 1)

    with ops.Graph().as_default():
      imported = load.load(path)
      keys = constant_op.constant(["brain", "test", "foo", "surgery"])
      output1 = imported.lookup1(keys)
      output2 = imported.lookup2(keys)
      with monitored_session.MonitoredSession() as sess:
        self.assertAllEqual([0, -1, -1, 2], sess.run(output1))
        self.assertAllEqual([2, 0, 1, -1], sess.run(output2))

  def test_preserve_argspec(self, cycles):

    def f(a, b, c):  # pylint: disable=unused-argument
      return None

    original_fullargspec = tf_inspect.getfullargspec(f)

    root = tracking.AutoTrackable()
    root.f = def_function.function(f)
    imported = cycle(root, cycles)

    restored_fullargspec = tf_inspect.getfullargspec(imported.f)
    self.assertEqual(original_fullargspec, restored_fullargspec)

  def test_canonicalize_inputs(self, cycles):
    @def_function.function(autograph=False)
    def func(a=1, b=2, c=3, training=True):
      if training:
        return [a, b, c, training]
      else:
        return [c, b, a, training]

    # TODO(b/123501567): Work-around to trigger generic traces of a function
    # with extra non tensor args.
    signature = 3*[tensor_spec.TensorSpec(None, dtypes.float32)]
    @def_function.function(input_signature=signature)
    def trigger(a, b, c):
      func(a, b, c, True)
      func(a, b, c, False)

    trigger.get_concrete_function()

    root = tracking.AutoTrackable()
    root.f = func
    root = cycle(root, cycles)
    self.assertAllEqual(root.f(), [1.0, 2.0, 3.0, True])
    self.assertAllEqual(root.f(-1.0, training=False), [3.0, 2.0, -1.0, False])

    with self.assertRaisesRegex(ValueError, "Could not find matching function"):
      root.f(["hello", 1.0])

  def test_prefer_specific_trace(self, cycles):
    @def_function.function(autograph=False)
    def func(a):
      if isinstance(a, int):
        return a
      else:
        return a + 1

    self.assertAllEqual(2, func(2).numpy())
    self.assertAllEqual(3, func(constant_op.constant(2)).numpy())

    root = tracking.AutoTrackable()
    root.f = func
    root = cycle(root, cycles)
    self.assertAllEqual(2, root.f(2).numpy())
    self.assertAllEqual(4, root.f(3).numpy())
    self.assertAllEqual(3, root.f(constant_op.constant(2)).numpy())
    self.assertAllEqual(4, root.f(constant_op.constant(3)).numpy())

  def test_partial(self, cycles):
    def f(x, y):
      return x + y

    func = def_function.function(
        functools.partial(f, x=array_ops.zeros([1]), y=array_ops.ones([1])))

    root = tracking.AutoTrackable()
    root.f = func
    self.assertAllEqual(root.f(), [1.0])

    root = cycle(root, cycles)
    self.assertAllEqual(root.f(), [1.0])

  def test_partial_with_non_tensor_defaults(self, cycles):

    def f(x, y=3):
      return x + y

    func = def_function.function(functools.partial(f, y=5))

    root = tracking.AutoTrackable()
    root.f = func
    self.assertAllEqual(root.f(1), 6)

    root = cycle(root, cycles)
    self.assertAllEqual(root.f(1), 6)

  def test_partial_with_positional(self, cycles):
    def f(x, y):
      return x + y

    func = def_function.function(functools.partial(f, constant_op.constant(5)))

    root = tracking.AutoTrackable()
    root.f = func
    self.assertAllEqual(root.f(1), 6)

    root = cycle(root, cycles)
    self.assertAllEqual(root.f(1), 6)

  def test_partial_with_positional_captured_tensors(self, cycles):

    def f(x, y):
      return x + y

    tensor = constant_op.constant(5) + constant_op.constant(7)
    func = def_function.function(functools.partial(f, tensor))

    root = tracking.AutoTrackable()
    root.f = func
    self.assertAllEqual(root.f(1), 13)

    root = cycle(root, cycles)
    self.assertAllEqual(root.f(1), 13)

  def test_partial_keyword_hiding_default(self, cycles):

    def f(x=3, training=True, y=7):
      if training:
        return x + y
      else:
        return x + y + 2

    func = def_function.function(functools.partial(f, y=6))

    root = tracking.AutoTrackable()
    root.f = func
    self.assertEqual(root.f().numpy(), 9)
    self.assertEqual(root.f(training=False).numpy(), 11)

    root = cycle(root, cycles)
    self.assertEqual(root.f().numpy(), 9)
    self.assertEqual(root.f(training=False).numpy(), 11)

  def test_partial_with_kwargs(self, cycles):

    def f(a, b, *args, **kwargs):
      args_sum = sum(args)
      return a + b + kwargs["some_tensor"] * kwargs["learning_rate"] + args_sum

    constant_tensor = constant_op.constant(10)
    func = def_function.function(
        functools.partial(
            f, 7, 1, 2, learning_rate=3, some_tensor=constant_tensor))

    root = tracking.AutoTrackable()
    root.f = func
    self.assertEqual(root.f(constant_op.constant(4)).numpy(), 44)

    root = cycle(root, cycles)
    self.assertEqual(root.f(constant_op.constant(5)).numpy(), 45)

  def test_partial_bind_only_first_argument(self, cycles):
    if sys.version_info[0] < 3:
      self.skipTest("Test is only valid in python3. Only then we get some more "
                    "advanced inspection of partials where this is allowed.")

    def f(x, y):
      return x + y

    partial_func = functools.partial(f, x=5)
    tf_func = def_function.function(partial_func)

    root = tracking.AutoTrackable()
    root.f = tf_func
    self.assertAllEqual(root.f(y=constant_op.constant(7)), 12)

    root = cycle(root, cycles)
    self.assertAllEqual(root.f(y=constant_op.constant(9)), 14)

  def test_partial_with_passed_fn_as_default(self, cycles):

    def f(x, y):
      return x(3) + y

    def my_func(a):
      return 2 * a

    func = def_function.function(functools.partial(f, my_func))

    root = tracking.AutoTrackable()
    root.f = func
    self.assertEqual(root.f(constant_op.constant(3)).numpy(), 9)

    root = cycle(root, cycles)
    self.assertEqual(root.f(constant_op.constant(3)).numpy(), 9)

  def test_partial_with_input_signature(self, cycles):

    def full_function(a, b, c=3.0):
      return a, b, c

    partial = functools.partial(full_function, 1, c=4)
    self.assertAllEqual((1, 2.0, 4), partial(2.0))

    signature = [tensor_spec.TensorSpec([], dtypes.float32)]
    func = def_function.function(partial, input_signature=signature)

    root = tracking.AutoTrackable()
    root.f = func
    a, b, c = root.f(2.0)
    self.assertAllEqual([a.numpy(), b.numpy(), c.numpy()], (1, 2.0, 4))

    root = cycle(root, cycles)
    a, b, c = root.f(3.0)
    self.assertAllEqual([a.numpy(), b.numpy(), c.numpy()], (1, 3.0, 4))

  def test_convert_to_input_signature(self, cycles):

    @def_function.function(
        input_signature=[tensor_spec.TensorSpec([None], dtypes.int32)])
    def func(x):
      return x

    root = tracking.AutoTrackable()
    root.f = func

    root = cycle(root, cycles)

    self.assertEqual([2], root.f([2]).numpy())

  def test_named_tuple(self, cycles):

    class NamedTupleType(collections.namedtuple("NamedTupleType", ["a", "b"])):
      pass

    @def_function.function
    def f(x):
      return x.a + x.b

    f.get_concrete_function(
        NamedTupleType(
            a=tensor_spec.TensorSpec(None, dtypes.float32, name="a"),
            b=tensor_spec.TensorSpec(None, dtypes.float32, name="b")))
    obj = tracking.AutoTrackable()
    obj.__call__ = f
    if sys.version_info.major == 3 and sys.version_info.minor < 5:
      # TODO(allenl): figure out why this doesn't work in Python3.4
      self.skipTest("Not working in Python 3.4")
    imported = cycle(obj, cycles)
    self.assertAllClose(3.,
                        imported(NamedTupleType(a=constant_op.constant(1.),
                                                b=constant_op.constant(2.))))

  def test_extra_args(self, cycles):

    @def_function.function
    def f(x):
      return math_ops.add(x["a"], 1.)
    # Trigger a trace.
    f({"a": constant_op.constant(2.0)})

    obj = tracking.AutoTrackable()
    obj.__call__ = f
    imported = cycle(obj, cycles)

    self.assertEqual(4.0, imported({"a": 3.0}).numpy())

    with self.assertRaisesRegex(ValueError,
                                "Could not find matching function to call"):
      imported({"a": 2.0, "b": 3.0})

  def test_shapes_available(self, cycles):

    @def_function.function(input_signature=[
        tensor_spec.TensorSpec([None, 3], dtypes.int32),
        tensor_spec.TensorSpec([None, 2], dtypes.int32)
    ])
    def func(x, y):
      return array_ops.concat([x, y], axis=1)

    root = tracking.AutoTrackable()
    root.f = func

    root = cycle(root, cycles)

    imported_graph = root.f.get_concrete_function().graph
    input_x, input_y = imported_graph.inputs
    self.assertEqual([None, 3], input_x.shape.as_list())
    self.assertEqual([None, 2], input_y.shape.as_list())
    output, = imported_graph.outputs
    self.assertEqual([None, 5], output.shape.as_list())
    signature = root.signatures["serving_default"]
    self.assertEqual(
        [None, 3], signature.inputs[0].shape.as_list())
    self.assertEqual(
        [None, 2], signature.inputs[1].shape.as_list())
    self.assertEqual(
        [None, 5], signature.outputs[0].shape.as_list())

  def test_variables_destroyed(self, cycles):
    v1 = variables.Variable(1.)
    weak_v1 = weakref.ref(v1)
    root = util.Checkpoint(v=v1)
    root = cycle(root, cycles)
    del v1
    self.assertIsNone(weak_v1())
    weak_v2 = weakref.ref(root.v)
    del root
    self.assertIsNone(weak_v2())

  def test_variable_attributes_preserved(self, cycles):
    v = variables.Variable(
        1.,
        trainable=False,
        synchronization=variables.VariableSynchronization.NONE,
        aggregation=variables.VariableAggregation.ONLY_FIRST_REPLICA)
    self.assertEqual(variables.VariableSynchronization.NONE,
                     v.synchronization)
    self.assertEqual(variables.VariableAggregation.ONLY_FIRST_REPLICA,
                     v.aggregation)
    root = tracking.AutoTrackable()
    root.v = v
    root = cycle(root, cycles)
    self.assertEqual(False, root.v.trainable)
    self.assertEqual(variables.VariableSynchronization.NONE,
                     root.v.synchronization)
    self.assertEqual(variables.VariableAggregation.ONLY_FIRST_REPLICA,
                     root.v.aggregation)

  def test_captured_dataset(self, cycles):

    class HasDataset(module.Module):

      def __init__(self):
        super(HasDataset, self).__init__()
        self.dataset = (
            dataset_ops.Dataset.range(5)
            .map(lambda x: x ** 2))

      @def_function.function
      def __call__(self, x):
        current_sum = array_ops.zeros([], dtype=dtypes.int64)
        for element in self.dataset:
          current_sum += x * element
        return current_sum

    root = HasDataset()
    self.assertEqual(
        3 * (1 + 4 + 9 + 16),
        root(constant_op.constant(3, dtype=dtypes.int64)).numpy())
    root = cycle(root, cycles)
    self.assertEqual(
        3 * (1 + 4 + 9 + 16),
        root(constant_op.constant(3, dtype=dtypes.int64)).numpy())

  def test_tuple_signature(self, cycles):
    root = util.Checkpoint()
    root.f = def_function.function(
        lambda: (array_ops.ones([]), array_ops.zeros([])),
        input_signature=())
    for _ in range(cycles):
      root = cycle(root, 1, signatures=root.f)
    self.assertEqual(({"output_0": 1., "output_1": 0.}),
                     self.evaluate(root.signatures["serving_default"]()))

  def test_version_info(self, cycles):
    root = util.Checkpoint()
    root = cycle(root, cycles)
    self.assertEqual(versions.__version__, root.tensorflow_version)
    self.assertEqual(versions.__git_version__, root.tensorflow_git_version)

  def test_load_grad_save(self, cycles):
    root = util.Checkpoint()
    root.v = variables.Variable(2.)
    root.f = def_function.function(lambda x: root.v * x)
    root.g = def_function.function(root.f)
    for _ in range(cycles):
      with backprop.GradientTape() as tape:
        inp = constant_op.constant(2.)
        tape.watch(inp)
        output = root.g(inp)
        self.assertAllClose(4., output)
      self.assertAllClose(2., tape.gradient(output, inp))
      root = cycle(root, 1)

  def test_destroy_resource(self, cycles):

    def get_handle():
      return resource_variable_ops.var_handle_op(
          shape=tensor_shape.as_shape([]),
          dtype=dtypes.float32,
          shared_name="my_var_name",
          name="my_var",
          container="my_container")

    class MyResourceDeleter(tracking.CapturableResourceDeleter):

      def destroy_resource(self):
        handle = get_handle()
        resource_variable_ops.destroy_resource_op(
            handle, ignore_lookup_error=True)

    class MyResource(tracking.TrackableResource):

      def __init__(self):
        # Set the resource deleter, so when the resource object goes out of
        # scope it will be deleted automatically.
        super(MyResource, self).__init__(deleter=MyResourceDeleter())

      def _create_resource(self):
        return get_handle()

      def _initialize(self):
        resource_variable_ops.assign_variable_op(
            self.resource_handle, 1.0, name="assign")

    class MyModel(tracking.AutoTrackable):

      def __init__(self):
        super(MyModel, self).__init__()
        self.resource = MyResource()

      @def_function.function(input_signature=[])
      def increase(self):
        handle = self.resource.resource_handle
        resource_variable_ops.assign_add_variable_op(
            handle, 10.0, name="assign_add")
        return resource_variable_ops.read_variable_op(handle, dtypes.float32)

    root = MyModel()
    imported = cycle(root, cycles)
    self.assertEqual(11, imported.increase().numpy())  # Create the resource.

    handle = imported.resource.resource_handle

    # Delete the imported SaveModel. Since we explicitly set the deleter, it
    # should destroy the resource automatically.
    del imported

    # Try to destroy the resource again, should fail.
    with self.assertRaisesRegex(errors.NotFoundError,
                                r"Resource .* does not exist."):
      resource_variable_ops.destroy_resource_op(
          handle, ignore_lookup_error=False)

  def test_function_called_as_operation(self, cycles):

    @framework_function.Defun(dtypes.float32)
    def inner(x):
      return x + 1.

    @def_function.function(
        input_signature=[tensor_spec.TensorSpec([], dtypes.float32)])
    def outer(x):
      return inner(x)

    root = module.Module()
    root.f = outer
    imported = cycle(root, cycles)
    self.assertAllClose(2., imported.f(constant_op.constant(1.)))

  def test_ragged(self, cycles):

    @def_function.function
    def f(x, c=1):
      """Returns Tensor x incremented by Python constant c."""
      return math_ops.add(x, c)

    for c in (1, 2, 3):
      _ = f.get_concrete_function(
          ragged_tensor.RaggedTensorSpec([None, None], dtype=dtypes.int32),
          c)

    obj = tracking.AutoTrackable()
    obj.f = f

    imported1 = cycle(obj, cycles, signatures={})
    rt = ragged_factory_ops.constant([[1, 2], [3]])
    self.assertAllEqual(imported1.f(rt), [[2, 3], [4]])
    self.assertAllEqual(imported1.f(rt, 2), [[3, 4], [5]])
    self.assertAllEqual(imported1.f(rt, 3), [[4, 5], [6]])

    imported2 = cycle(obj, cycles)
    rt = ragged_factory_ops.constant([[1, 2], [3]])
    self.assertAllEqual(imported2.f(rt, 1), [[2, 3], [4]])
    self.assertAllEqual(imported2.f(rt, 2), [[3, 4], [5]])
    self.assertAllEqual(imported2.f(rt, 3), [[4, 5], [6]])

  def test_accepts_io_device(self, cycles):
    options = load_options.LoadOptions()
    self.assertIsNone(options.experimental_io_device)
    options = load_options.LoadOptions(experimental_io_device="/job:localhost")
    self.assertEqual("/job:localhost", options.experimental_io_device)

  def test_load_custom_saveable_object(self, cycles):
    root = tracking.AutoTrackable()
    root.table = lookup_ops.MutableHashTable(dtypes.string, dtypes.float32, -1)
    root.table.insert("foo", 15)

    @def_function.function(
        input_signature=[tensor_spec.TensorSpec(None, dtypes.string)])
    def lookup(key):
      return root.table.lookup(key)

    root.lookup = lookup

    imported = cycle(root, cycles)
    self.assertEqual(self.evaluate(imported.lookup("foo")), 15)
    self.assertEqual(self.evaluate(imported.lookup("idk")), -1)

  def test_saving_ndarray_specs(self, cycles):
    class NdarrayModule(module.Module):

      @def_function.function
      def plain(self, x):
        return tnp.add(x, 1)

      @def_function.function(input_signature=[
          np_arrays.NdarraySpec(tensor_spec.TensorSpec([], dtypes.float32))])
      def with_signature(self, x):
        return tnp.add(x, 1)

    m = NdarrayModule()
    c = tnp.asarray(3.0, tnp.float32)
    output_plain, output_with_signature = m.plain(c), m.with_signature(c)

    loaded_m = cycle(m, cycles)

    load_output_plain, load_output_with_signature = (
        loaded_m.plain(c), loaded_m.with_signature(c))

    self.assertIsInstance(output_plain, tnp.ndarray)
    self.assertIsInstance(load_output_plain, tnp.ndarray)
    self.assertIsInstance(output_with_signature, tnp.ndarray)
    self.assertIsInstance(load_output_with_signature, tnp.ndarray)
    self.assertAllClose(output_plain, load_output_plain)
    self.assertAllClose(output_with_signature, load_output_with_signature)


class SingleCycleTests(test.TestCase, parameterized.TestCase):

  def test_load_with_tags(self):
    root = tracking.AutoTrackable()
    path = tempfile.mkdtemp(prefix=self.get_temp_dir())
    save.save(root, path)
    with self.assertRaises(ValueError):
      load.load(path, tags=[tag_constants.EVAL])
    load.load(path, tags=[tag_constants.SERVING])
    load.load(path, tags=tag_constants.SERVING)
    load.load(path, tags=set([tag_constants.SERVING]))

  def test_docstring_examples(self):
    path = tempfile.mkdtemp(prefix=self.get_temp_dir())
    exported = util.Checkpoint(v=variables.Variable(3.))
    exported.f = def_function.function(
        lambda x: exported.v * x,
        input_signature=[
            tensor_spec.TensorSpec(shape=None, dtype=dtypes.float32)])
    save.save(exported, path)
    imported = load.load(path)
    self.assertEqual(3., imported.v.numpy())
    self.assertEqual(6., imported.f(x=constant_op.constant(2.)).numpy())

    save.save(exported, path, exported.f.get_concrete_function())
    imported = load.load(path)
    f = imported.signatures["serving_default"]
    self.assertAllEqual(
        [[-3.]],
        f(x=constant_op.constant([[-1.]]))["output_0"].numpy())


  def test_object_with_extra_dependencies(self):

    class Extra(tracking.AutoTrackable):

      def _list_extra_dependencies_for_serialization(self, cache):
        if self not in cache:
          cache[self] = {"a": variables.Variable(5.)}
        return cache[self]
    root = Extra()
    path = tempfile.mkdtemp(prefix=self.get_temp_dir())
    save.save(root, path)
    imported = load.load(path)
    self.assertEqual(5, self.evaluate(imported.a))

    root.a = variables.Variable(3.)
    with self.assertRaisesRegex(
        ValueError, "object has an attribute named a, which is reserved."):
      save.save(root, path)

  def test_save_cached_variable(self):
    with ops.Graph().as_default(), session_lib.Session() as session:
      obj = tracking.AutoTrackable()
      obj.v = variables.Variable(2., caching_device=lambda op: op.device)
      obj.w = variables.Variable(3.)
      session.run([obj.v.initializer, obj.w.initializer])

      @def_function.function
      def total():
        return obj.v + obj.w

      @def_function.function(input_signature=[tensor_spec.TensorSpec([])])
      def wrapped_total(x):
        return total() + x

      @def_function.function
      def increment_v(x):
        obj.v.assign_add(x)

      session.run(increment_v(constant_op.constant(3.)))  # generate signatures
      self.assertAllClose(8, total())
      self.assertAllClose(13, wrapped_total(constant_op.constant(5.)))

      obj.total = total
      obj.wrapped_total = wrapped_total.get_concrete_function()
      obj.increment_v = increment_v

      save_dir = os.path.join(self.get_temp_dir(), "saved_model")
      save.save(obj, save_dir, signatures=total.get_concrete_function())
      imported = load.load(save_dir)
      session.run(variables.global_variables_initializer())
      self.assertAllClose(8, imported.total())
      session.run(imported.increment_v(4))
      self.assertAllClose(12, imported.total())
      self.assertAllClose(15, imported.wrapped_total(constant_op.constant(3.)))
      self.assertAllClose({"output_0": 12},
                          imported.signatures["serving_default"]())

    # Try loading and running the function in eager mode
    imported = load.load(save_dir)
    self.assertAllClose(8, imported.total())
    imported.increment_v(5)
    self.assertAllClose(13, imported.total())
    self.assertAllClose(13.5, imported.wrapped_total(constant_op.constant(.5)))
    self.assertAllClose({"output_0": 13},
                        imported.signatures["serving_default"]())

  # TODO(allenl, kkb): Use the new memory checker here once it's fast enough (3
  # iterations took hundreds of seconds). It would be really nice to check
  # allocations at a lower level.
  @test_util.assert_no_new_pyobjects_executing_eagerly
  def test_functions_cleaned(self):
    if sys.version_info.major < 3:
      self.skipTest("Not working in Python 2")
    root = module.Module()
    root.v = variables.Variable(1.)
    root.f = def_function.function(
        lambda x: x + root.v,
        input_signature=[
            tensor_spec.TensorSpec(shape=[], dtype=dtypes.float32)])
    cycle(root, 1)


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