# 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.
# ==============================================================================
import collections
import copy
import json
import os
import pickle

from absl.testing import parameterized

from tensorflow.python.checkpoint import checkpoint as util
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
from tensorflow.python.eager import test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import tensor_shape
from tensorflow.python.module import module
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import variables
from tensorflow.python.trackable import autotrackable
from tensorflow.python.trackable import data_structures
from tensorflow.python.util import nest
from tensorflow.python.util import serialization


class ListTests(test.TestCase):

  def testJSONSerialization(self):
    obj = autotrackable.AutoTrackable()
    obj.l = [1]
    json.dumps(obj.l, default=serialization.get_json_type)

  def testNotTrackable(self):
    class NotTrackable(object):
      pass

    with self.assertRaises(ValueError):
      data_structures.List([NotTrackable()])

  def testCallNotImplemented(self):
    with self.assertRaisesRegex(TypeError, "not callable"):
      data_structures.List()(1.)  # pylint: disable=not-callable

  def testNoPop(self):
    with self.assertRaises(AttributeError):
      data_structures.List().pop()

  def testNesting(self):
    with context.graph_mode():
      inner = data_structures.List()
      outer = data_structures.List([inner])

      class TestDense(module.Module):
        def __init__(self):
          self.w = variables.Variable(1.0)
          self.b = variables.Variable(0.0)

        def __call__(self, inputs):
          return self.w * inputs + self.b

      inner.append(TestDense())
      inner[0](array_ops.ones([2, 3]))
      self.assertLen(outer.variables, 2)
      self.assertIsInstance(
          outer.variables[0],
          resource_variable_ops.ResourceVariable)

  def testNonLayerVariables(self):
    v = resource_variable_ops.ResourceVariable([1.])
    l = data_structures.List([v])
    self.assertTrue(l.trainable)
    self.assertEqual([], l.layers)
    self.assertEqual([v], l.variables)
    self.assertEqual([v], l.trainable_weights)
    self.assertEqual([], l.non_trainable_variables)
    l.trainable = False
    self.assertEqual([v], l.variables)
    self.assertEqual([], l.trainable_variables)
    self.assertEqual([v], l.non_trainable_variables)
    l.trainable = True
    v2 = resource_variable_ops.ResourceVariable(1., trainable=False)
    l.append(v2)
    self.assertEqual([v, v2], l.weights)
    self.assertEqual([v], l.trainable_weights)
    self.assertEqual([v2], l.non_trainable_weights)

  def testCopy(self):
    v1 = resource_variable_ops.ResourceVariable(1.)
    v2 = resource_variable_ops.ResourceVariable(1.)
    v3 = resource_variable_ops.ResourceVariable(1.)

    l1 = data_structures.List([v1, v2])
    l2 = l1.copy()
    l2.append(v3)
    self.assertEqual(list(l1), [v1, v2])
    self.assertEqual(list(l2), [v1, v2, v3])

  def testSlicing(self):
    v1 = resource_variable_ops.ResourceVariable(1.)
    v2 = resource_variable_ops.ResourceVariable(1.)
    v3 = resource_variable_ops.ResourceVariable(1.)
    v4 = resource_variable_ops.ResourceVariable(1.)

    l = data_structures.List([v1, v2, v3, v4])
    self.assertEqual(l[1:], [v2, v3, v4])
    self.assertEqual(l[1:-1], [v2, v3])
    self.assertEqual(l[:-1], [v1, v2, v3])

  def testHash(self):
    has_sequences = {data_structures.List(), data_structures.List()}
    self.assertEqual(2, len(has_sequences))
    self.assertNotIn(data_structures.List(), has_sequences)

  def testIMul_zero(self):
    l = data_structures.List([])
    with self.assertRaisesRegex(ValueError, "List only supports append"):
      l *= 0

  def testIMul(self):
    v = resource_variable_ops.ResourceVariable(1.)
    l = data_structures.List([v])
    l *= 2
    self.assertEqual(list(l), [v] * 2)

  def testMul(self):
    v = resource_variable_ops.ResourceVariable(1.)
    l = data_structures.List([v, v, v])
    self.assertEqual(list(l * 2), [v, v, v] * 2)

  def testRMul(self):
    v = resource_variable_ops.ResourceVariable(1.)
    l = data_structures.List([v, v, v])
    self.assertEqual(list(2 * l), [v, v, v] * 2)


class ListWrapperTest(test.TestCase):

  IGNORED = ("__new__", "__init__", "__subclasshook__", "__getattribute__")

  def test_overrides_all_list_methods(self):
    not_overridden = []

    for name in dir(list):
      if name in ListWrapperTest.IGNORED:
        continue

      list_method = getattr(list, name)

      if not callable(list_method):
        continue

      object_method = getattr(object, name, None)
      if object_method is not None and object_method == list_method:
        # Skip methods that aren't overridden from object.
        continue

      if list_method == getattr(data_structures.ListWrapper, name):
        not_overridden.append(name)

    if not_overridden:
      self.fail("ListWrapper does not override %s" % (not_overridden))

  def testPickle(self):
    original = data_structures.ListWrapper([1, 2])
    serialized = pickle.dumps(original)
    del original
    deserialized = pickle.loads(serialized)
    self.assertEqual([1, 2], deserialized)

  def testSameStructure(self):
    l = [1]
    nest.assert_same_structure(l, data_structures.ListWrapper(copy.copy(l)))

  def testMutateWithoutTrackableComponents(self):
    m = module.Module()
    m.l = [1, 2]
    m.l.insert(0, 0)
    self.assertEqual(m.l, [0, 1, 2])
    self.assertEqual(m.l._trackable_children(), {})

  def testFunctionCaching(self):
    @def_function.function
    def f(list_input):
      return list_input[0] + constant_op.constant(1.)

    first_trace = f.get_concrete_function([constant_op.constant(2.)])
    second_trace = f.get_concrete_function(
        data_structures.ListWrapper([constant_op.constant(3.)]))
    self.assertIs(first_trace, second_trace)

  def testListWrapperBasic(self):
    # ListWrapper, unlike List, compares like the built-in list type (since it
    # is used to automatically replace lists).
    a = autotrackable.AutoTrackable()
    b = autotrackable.AutoTrackable()
    self.assertEqual([a, a],
                     [a, a])
    self.assertEqual(data_structures.ListWrapper([a, a]),
                     data_structures.ListWrapper([a, a]))
    self.assertEqual([a, a],
                     data_structures.ListWrapper([a, a]))
    self.assertEqual(data_structures.ListWrapper([a, a]),
                     [a, a])
    self.assertNotEqual([a, a],
                        [b, a])
    self.assertNotEqual(data_structures.ListWrapper([a, a]),
                        data_structures.ListWrapper([b, a]))
    self.assertNotEqual([a, a],
                        data_structures.ListWrapper([b, a]))
    self.assertLess([a], [a, b])
    self.assertLess(data_structures.ListWrapper([a]),
                    data_structures.ListWrapper([a, b]))
    self.assertLessEqual([a], [a, b])
    self.assertLessEqual(data_structures.ListWrapper([a]),
                         data_structures.ListWrapper([a, b]))
    self.assertGreater([a, b], [a])
    self.assertGreater(data_structures.ListWrapper([a, b]),
                       data_structures.ListWrapper([a]))
    self.assertGreaterEqual([a, b], [a])
    self.assertGreaterEqual(data_structures.ListWrapper([a, b]),
                            data_structures.ListWrapper([a]))
    self.assertEqual([a], data_structures.ListWrapper([a]))
    self.assertEqual([a], list(data_structures.List([a])))
    self.assertEqual([a, a], data_structures.ListWrapper([a]) + [a])
    self.assertEqual([a, a], [a] + data_structures.ListWrapper([a]))
    self.assertIsInstance(data_structures.ListWrapper([a]), list)
    self.assertEqual(
        tensor_shape.TensorShape([None, 2]).as_list(),
        (data_structures.ListWrapper([None])
         + tensor_shape.TensorShape([2])).as_list())

  def testAcceptsNonTrackableContent(self):
    l = data_structures.ListWrapper([1, 2, 3])
    self.assertEqual(l, [1, 2, 3])

  def testWrapperChangesList(self):
    l = []
    l_wrapper = data_structures.ListWrapper(l)
    l_wrapper.append(1)
    self.assertEqual([1], l)

  def testListChangesWrapper(self):
    l = []
    l_wrapper = data_structures.ListWrapper(l)
    l.append(1)
    self.assertEqual([1], l_wrapper)

  def testNotHashable(self):
    with self.assertRaises(TypeError):
      hash(data_structures.ListWrapper())  # pylint: disable=no-value-for-parameter

  def testDelItem(self):
    l = data_structures.ListWrapper([1, 2, 3, [4]])
    del l[0]
    self.assertEqual(l, [2, 3, [4]])
    self.assertUnableToSave(l, "Unable to save .*__delitem__")

  def testDelSlice(self):
    l = data_structures.ListWrapper([1, 2, 3, [4]])
    del l[2:3]
    self.assertEqual(l, [1, 2, [4]])
    self.assertUnableToSave(l, "Unable to save .*__delslice__")

  def testSetSlice_canSaveForNonTrackableItems(self):
    l = data_structures.ListWrapper([1, 2, 3, 4])
    l[:] = 2, 8, 9, 0
    self.assertEqual(l, [2, 8, 9, 0])
    l._maybe_initialize_trackable()  # pylint: disable=protected-access
    self.assertEqual(len(l._trackable_children()), 0)  # pylint: disable=protected-access

  def testSetSlice_cannotSaveIfTrackableModified(self):
    v1 = resource_variable_ops.ResourceVariable(1.)
    v2 = resource_variable_ops.ResourceVariable(1.)
    l = data_structures.ListWrapper([1, 2, v1, v2])
    l[:] = 2, 8, 9, v2
    self.assertEqual(l, [2, 8, 9, v2])
    self.assertUnableToSave(l, "Unable to save .*__setslice__")

  def testSetSlice_truncate(self):
    l = data_structures.ListWrapper([1, 2, 3, 4])
    l[:] = []
    self.assertEqual(l, [])

  def testSetSlice_extend(self):
    l = data_structures.ListWrapper([1, 2, 3, 4])
    l[2:] = 1, 2, 3, 4
    self.assertEqual(l, [1, 2, 1, 2, 3, 4])

  def testIMulNegative(self):
    l = data_structures.ListWrapper([1, 2, 3, [4]])
    l *= -1
    self.assertEqual(l, [1, 2, 3, [4]] * -1)
    self.assertUnableToSave(l, "Unable to save")

  def testIMulPositive(self):
    v = variables.Variable(1.)
    l = data_structures.ListWrapper([1, 2, 3, 4, v])
    self.assertDictEqual({"4": v}, l._trackable_children())
    root = util.Checkpoint(l=l)
    prefix = os.path.join(self.get_temp_dir(), "ckpt")
    path = root.save(prefix)
    v.assign(5.)
    l *= 2
    self.assertEqual(l, [1, 2, 3, 4, v, 1, 2, 3, 4, v])
    self.assertDictEqual({"4": v, "9": v}, l._trackable_children())
    root.restore(path)
    self.assertAllClose(1., v.numpy())

  def testSort(self):
    l = data_structures.ListWrapper([[1], [2], [3], [4]])
    l.sort()
    self.assertAllEqual(l, [[1], [2], [3], [4]])
    # Regardless of being a no-op for the input list, we still refuse to save.
    # This is intentional since otherwise we would end up with a hard to debug
    # case for users (e.g. sometimes sort on a ListWrapper is trackable and
    # other times it is not).
    self.assertUnableToSave(l, "Unable to save .*sort")

  def assertUnableToSave(self, l, msg):
    l._maybe_initialize_trackable()  # pylint: disable=protected-access
    with self.assertRaisesRegex(ValueError, msg):
      return l._trackable_children()  # pylint: disable=protected-access


class MappingTests(test.TestCase):

  def testJSONSerialization(self):
    obj = autotrackable.AutoTrackable()
    obj.d = {"a": 2}
    json.dumps(obj.d, default=serialization.get_json_type)

  def testNoOverwrite(self):
    mapping = data_structures.Mapping()
    original = data_structures.List()
    mapping["a"] = original
    with self.assertRaises(ValueError):
      mapping["a"] = data_structures.List()
    self.assertIs(original, mapping["a"])
    with self.assertRaises(AttributeError):
      del mapping["a"]  # pylint: disable=unsupported-delete-operation
    mapping.update(b=data_structures.Mapping())
    with self.assertRaises(ValueError):
      mapping.update({"b": data_structures.Mapping()})

  def testNonStringKeys(self):
    mapping = data_structures.Mapping()
    with self.assertRaises(TypeError):
      mapping[1] = data_structures.List()

  def testHashing(self):
    has_mappings = set([data_structures.Mapping(),
                        data_structures.Mapping()])
    self.assertEqual(2, len(has_mappings))
    self.assertNotIn(data_structures.Mapping(), has_mappings)
    # In contrast to Mapping, dict wrappers are not hashable
    a = autotrackable.AutoTrackable()
    a.d = {}
    self.assertEqual({}, a.d)
    self.assertFalse({} != a.d)  # pylint: disable=g-explicit-bool-comparison
    self.assertNotEqual({1: 2}, a.d)
    with self.assertRaisesRegex(TypeError, "unhashable"):
      set([a.d])

  def testListShallowCopy(self):
    root = autotrackable.AutoTrackable()
    orig_list = [[1.]]
    root.a = orig_list
    copied = copy.copy(root.a)
    self.assertAllEqual([[1.]], copied)
    self.assertIsNot(root.a, copied)
    self.assertIs(root.a[0], copied[0])

    # Dirtiness should be inherited
    util.list_objects(root.a)
    orig_list.append(1.)
    with self.assertRaises(ValueError):
      util.list_objects(root.a)
    with self.assertRaises(ValueError):
      util.list_objects(copy.copy(root.a))

  def testListDeepCopy(self):
    root = autotrackable.AutoTrackable()
    orig_list = [[1.]]
    root.a = orig_list
    copied = copy.deepcopy(root.a)
    self.assertAllEqual([[1.]], copied)
    self.assertIsNot(root.a, copied)
    self.assertIsNot(root.a[0], copied[0])

    # Dirtiness should be inherited
    util.list_objects(root.a)
    orig_list.append(1.)
    with self.assertRaises(ValueError):
      util.list_objects(root.a)
    with self.assertRaises(ValueError):
      util.list_objects(copy.deepcopy(root.a))

  def testDictShallowCopy(self):
    root = autotrackable.AutoTrackable()
    orig_dict = {"a": [1.]}
    root.a = orig_dict
    copied = copy.copy(root.a)
    self.assertAllEqual([1.], copied["a"])
    self.assertIsNot(root.a, copied)
    self.assertIs(root.a["a"], copied["a"])

    copied = root.a.copy()
    self.assertAllEqual([1.], copied["a"])
    self.assertIsNot(root.a, copied)
    self.assertIs(root.a["a"], copied["a"])

    # Dirtiness should be inherited
    util.list_objects(root.a)
    orig_dict["b"] = []
    with self.assertRaises(ValueError):
      util.list_objects(root.a)
    with self.assertRaises(ValueError):
      util.list_objects(copy.copy(root.a))

  def testDictDeepCopy(self):
    root = autotrackable.AutoTrackable()
    orig_dict = {"a": [1.]}
    root.a = orig_dict
    copied = copy.deepcopy(root.a)
    self.assertAllEqual([1.], copied["a"])
    self.assertIsNot(root.a, copied)
    self.assertIsNot(root.a["a"], copied["a"])

    # Dirtiness should be inherited
    util.list_objects(root.a)
    orig_dict["b"] = []
    with self.assertRaises(ValueError):
      util.list_objects(root.a)
    with self.assertRaises(ValueError):
      util.list_objects(copy.deepcopy(root.a))

  def testShallowCopyTrackable(self):
    original = autotrackable.AutoTrackable()
    original_sub = autotrackable.AutoTrackable()
    original.a = [[1.]]
    original.b = {"a": original_sub}
    shallow_copied = copy.copy(original)
    self.assertIs(original_sub, shallow_copied.b["a"])
    self.assertIsNot(original, shallow_copied)
    self.assertEqual([[1.]], shallow_copied.a)
    shallow_deps = util.list_objects(shallow_copied)
    self.assertIn(shallow_copied.a, shallow_deps)
    self.assertIn(shallow_copied.b, shallow_deps)
    self.assertIn(shallow_copied.b["a"], shallow_deps)

  def testDeepCopyTrackable(self):
    original = autotrackable.AutoTrackable()
    original_sub = autotrackable.AutoTrackable()
    original.a = [[1.]]
    original.b = {"a": original_sub}
    self.assertIsInstance(original.b, dict)
    deep_copied = copy.deepcopy(original)
    self.assertIsInstance(deep_copied.b, dict)
    self.assertIsNot(original, deep_copied)
    self.assertIsNot(original_sub, deep_copied.b["a"])
    self.assertEqual([[1.]], deep_copied.a)
    self.assertIsInstance(deep_copied.b["a"], autotrackable.AutoTrackable)
    deps = util.list_objects(deep_copied)
    self.assertIn(deep_copied.a, deps)
    self.assertIn(deep_copied.b, deps)
    self.assertIn(deep_copied.b["a"], deps)
    self.assertNotIn(original_sub, deps)

  def testConstructableFromSequence(self):
    result = data_structures._DictWrapper([(1, 2), (3, 4)])
    self.assertIsInstance(result, dict)
    self.assertEqual({1: 2, 3: 4}, result)

  def testPickle(self):
    original = data_structures._DictWrapper(dict(a=1, b=2))
    serialized = pickle.dumps(original)
    del original
    deserialized = pickle.loads(serialized)
    self.assertEqual(dict(a=1, b=2), deserialized)

  def testListAddOrder(self):
    self.assertEqual([1., 2.],
                     data_structures.ListWrapper([1.])
                     + data_structures.ListWrapper([2.]))
    self.assertEqual([1., 2.],
                     data_structures.ListWrapper([1.])
                     + [2.])
    self.assertEqual([1., 2.],
                     [1.]
                     + data_structures.ListWrapper([2.]))

  def testSameStructure(self):
    d = {1: "a"}
    nest.assert_same_structure(d, data_structures._DictWrapper(d.copy()))

  def testFunctionCaching(self):
    @def_function.function
    def f(dict_input):
      return dict_input["x"] + constant_op.constant(1.)

    first_trace = f.get_concrete_function({"x": constant_op.constant(2.)})
    second_trace = f.get_concrete_function(
        data_structures._DictWrapper({"x": constant_op.constant(3.)}))
    self.assertIs(first_trace, second_trace)


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

  def testJSONSerialization(self):
    obj = autotrackable.AutoTrackable()
    obj.l = (1,)
    json.dumps(obj.l, default=serialization.get_json_type)

  def testNonLayerVariables(self):
    v = resource_variable_ops.ResourceVariable([1.])
    l = data_structures._TupleWrapper((v,))
    self.assertEqual([], l.layers)
    self.assertEqual([v], l.variables)
    self.assertEqual([v], l.trainable_weights)
    self.assertEqual([], l.non_trainable_variables)

  def testCopy(self):
    v1 = resource_variable_ops.ResourceVariable(1.)
    v2 = resource_variable_ops.ResourceVariable(1.)

    l1 = data_structures._TupleWrapper((v1, v2))
    l2 = copy.copy(l1)
    self.assertEqual(l1, (v1, v2))
    self.assertEqual(l2, (v1, v2))
    self.assertIs(l1[0], l2[0])
    l2_deep = copy.deepcopy(l1)
    self.assertIsNot(l1[0], l2_deep[0])
    with self.assertRaises(AttributeError):
      l2.append(v1)

  def testSlicing(self):
    v1 = resource_variable_ops.ResourceVariable(1.)
    v2 = resource_variable_ops.ResourceVariable(1.)
    v3 = resource_variable_ops.ResourceVariable(1.)
    v4 = resource_variable_ops.ResourceVariable(1.)

    l = data_structures._TupleWrapper((v1, v2, v3, v4))
    self.assertEqual(l[1:], (v2, v3, v4))
    self.assertEqual(l[1:-1], (v2, v3))
    self.assertEqual(l[:-1], (v1, v2, v3))

  def testHash(self):
    has_sequences = set([data_structures._TupleWrapper(),
                         data_structures._TupleWrapper()])
    self.assertLen(has_sequences, 1)
    self.assertIn(data_structures._TupleWrapper(), has_sequences)

  def testIMul_zero(self):
    l = data_structures._TupleWrapper((1,))
    l *= 0
    self.assertEqual((), l)

  def testIMul(self):
    # Note: tuple behavior differs from list behavior. Lists are mutated by
    # imul/iadd, tuples assign a new object to the left hand side of the
    # expression.
    v = resource_variable_ops.ResourceVariable(1.)
    l = data_structures._TupleWrapper((v,))
    original = l
    l *= 2
    self.assertEqual(l, (v,) * 2)
    self.assertNotEqual(original, (v,) * 2)

  def testIAdd(self):
    v = resource_variable_ops.ResourceVariable(1.)
    l = data_structures._TupleWrapper((v,))
    original = l
    l += (1,)
    self.assertEqual(l, (v, 1))
    self.assertNotEqual(original, (v, 1))
    self.assertEqual(original, (v,))

  def testMul(self):
    v = resource_variable_ops.ResourceVariable(1.)
    l = data_structures._TupleWrapper((v, v, v))
    self.assertEqual(l * 2, (v, v, v) * 2)

  def testRMul(self):
    v = resource_variable_ops.ResourceVariable(1.)
    l = data_structures._TupleWrapper((v, v, v))
    self.assertEqual(2 * l, (v, v, v) * 2)

  def testPickle(self):
    original = data_structures._TupleWrapper((1, 2))
    serialized = pickle.dumps(original)
    del original
    deserialized = pickle.loads(serialized)
    self.assertEqual((1, 2), deserialized)

  def testNamedTuple(self):
    named = collections.namedtuple("Named", ("x", "y"))
    v = variables.Variable(2)
    nt = named(x=v, y=2)
    m = module.Module()
    m.nt = nt
    self.assertIs(v, m.nt.x)
    self.assertIs(v, m.nt[0])
    self.assertIs(
        v, m._trackable_children()["nt"]._trackable_children()["x"])
    self.assertEqual(2, m.nt.y)

  def testNamedTupleConflictingAttributes(self):
    named = collections.namedtuple("Named", ("x", "weights"))
    v = variables.Variable(2)
    nt = named(x=v, weights=3)
    m = module.Module()
    m.nt = nt
    self.assertEqual(3, m.nt.weights)

  def testNamedSubclassing(self):
    named = collections.namedtuple("Named", ("x", "y"))
    v = variables.Variable(2)

    class NamedSubclass(named):

      def __new__(cls, x, y):
        del y  # unused
        return super(NamedSubclass, cls).__new__(cls, x, 3)

      @property
      def summed(self):
        return self.x + self.y

    nt = NamedSubclass(x=v, y=2)
    m = module.Module()
    m.nt = nt
    self.assertEqual(3, m.nt.y)
    self.assertIs(v, m.nt.x)
    self.assertIn(v,
                  m._trackable_children()["nt"]._trackable_children().values())
    self.assertIn("x", m.nt._trackable_children())
    self.assertIn("0", m.nt._trackable_children())
    self.assertEqual(5, self.evaluate(m.nt.summed))

  def testUnnamedSubclassing(self):
    v = variables.Variable(2)

    class UnnamedSubclass(tuple):

      @property
      def summed(self):
        return self[0] + self[1]

    unt = UnnamedSubclass([v, 2])
    m = module.Module()
    m.unt = unt
    self.assertIn("0", m.unt._trackable_children())
    self.assertLen(m.unt._trackable_children(), 1)
    self.assertEqual(4, self.evaluate(m.unt.summed))
    nest.assert_same_structure(
        [m.unt], nest.map_structure(lambda x: x, [m.unt]))

  def testNamedtupleSubclassWithCustomNew(self):
    class SubclassWithDifferentArgs(collections.namedtuple("A", ["x"])):

      def __new__(cls):
        return super(SubclassWithDifferentArgs, cls).__new__(cls, [])

    nt = SubclassWithDifferentArgs()
    m = module.Module()
    m.nt = nt
    m.nt.x.append(variables.Variable(1.))
    prefix = os.path.join(self.get_temp_dir(), "ckpt")
    ckpt = util.Checkpoint(m=m)
    with self.assertRaises(ValueError):
      ckpt.save(prefix)

  def testSameStructure(self):
    t = (variables.Variable(1.),)
    m = module.Module()
    m.t = t
    nest.assert_same_structure(t, m.t)
    nest.assert_same_structure(m.t, t)

    nt_type = collections.namedtuple("nt", ["x", "y"])
    nt = nt_type(x=1, y=2)
    m.nt = nt
    nest.assert_same_structure(m.nt, nt)
    with self.assertRaises(TypeError):  # pylint: disable=g-error-prone-assert-raises
      nest.assert_same_structure(m.nt, m.t)

  def testFlatten(self):
    t = data_structures._TupleWrapper((1, data_structures._TupleWrapper((2,))))
    self.assertEqual([1, 2], nest.flatten(t))
    self.assertEqual(
        nest.flatten_with_tuple_paths((1, (2,))),
        nest.flatten_with_tuple_paths(t))
    self.assertEqual((3, (4,)),
                     nest.pack_sequence_as(t, [3, 4]))
    nt_type = collections.namedtuple("nt", ["x", "y"])
    nt = nt_type(1., 2.)
    wrapped_nt = data_structures._TupleWrapper(nt)
    self.assertEqual(
        nest.flatten_with_tuple_paths(nt),
        nest.flatten_with_tuple_paths(wrapped_nt))
    self.assertEqual((3, 4,),
                     nest.pack_sequence_as(wrapped_nt, [3, 4]))
    self.assertEqual(3, nest.pack_sequence_as(wrapped_nt, [3, 4]).x)

  def testFunctionCaching(self):
    @def_function.function
    def f(tuple_input):
      return tuple_input[0] + constant_op.constant(1.)

    first_trace = f.get_concrete_function((constant_op.constant(2.),))
    second_trace = f.get_concrete_function(
        data_structures._TupleWrapper((constant_op.constant(3.),)))
    self.assertIs(first_trace, second_trace)

  def testPythonMapImpl(self):
    t = data_structures._TupleWrapper((1, data_structures._TupleWrapper((2,))))
    self.assertEqual(
        (4, (5,)),
        nest.map_structure_up_to((None, (None,)), lambda x: x + 3, t,
                                 check_types=True))
    nest.assert_shallow_structure((None, None), t)

  def testDatasetMap(self):
    dataset = dataset_ops.Dataset.from_tensor_slices(
        constant_op.constant([1, 2, 3]))
    dataset = dataset.map(lambda x: data_structures._TupleWrapper((x,)))
    for index, element in enumerate(dataset):
      self.assertEqual((index + 1,), self.evaluate(element))

  def testDatasetMapNamed(self):
    nt_type = collections.namedtuple("A", ["x"])
    dataset = dataset_ops.Dataset.from_tensor_slices(
        constant_op.constant([1, 2, 3]))
    dataset = dataset.map(lambda x: data_structures._TupleWrapper(nt_type(x)))
    for index, element in enumerate(dataset):
      self.assertEqual((index + 1,), self.evaluate(element))

  def testLoopAssignedModule(self):
    m = module.Module()
    m.s = (m,)
    self.assertLen(m._trackable_children(), 1)
    self.assertIn("s", m._trackable_children())
    self.assertIs(m.s, m._trackable_children()["s"])
    self.assertEqual((), m.trainable_variables)


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