# 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 save."""

import tensorflow.compat.v2 as tf

import os
from tensorflow.python.framework import test_util
from keras.layers import core
from keras.optimizer_v2 import adam


class _ModelWithOptimizerUsingDefun(tf.train.Checkpoint):

  def __init__(self):
    self.dense = core.Dense(1)
    self.optimizer = adam.Adam(0.01)

  @tf.function(
      input_signature=(tf.TensorSpec([None, 2], tf.float32),
                       tf.TensorSpec([None], tf.float32)),
  )
  def call(self, x, y):
    with tf.GradientTape() as tape:
      loss = tf.reduce_mean((self.dense(x) - y) ** 2.)
    trainable_variables = self.dense.trainable_variables
    gradients = tape.gradient(loss, trainable_variables)
    self.optimizer.apply_gradients(zip(gradients, trainable_variables))
    return {"loss": loss}


class MemoryTests(tf.test.TestCase):

  def setUp(self):
    super(MemoryTests, self).setUp()
    self._model = _ModelWithOptimizerUsingDefun()

  @test_util.assert_no_garbage_created
  def test_no_reference_cycles(self):
    x = tf.constant([[3., 4.]])
    y = tf.constant([2.])
    self._model.call(x, y)
    save_dir = os.path.join(self.get_temp_dir(), "saved_model")
    tf.saved_model.save(self._model, save_dir, self._model.call)


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