code: # 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.
  # ==============================================================================
  # pylint: disable=protected-access
  """Tests for saving and loading Keras models and layers from SavedModel.
  
  These should ensure that all layer properties are correctly assigned after
  loading from the SavedModel.
  
  Tests that focus on the model structure should go in revive_structure_test.py
  """
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import tensorflow as tf
  
  import os
  import shutil
  
  from absl.testing import parameterized
  import numpy as np
  
  from tensorflow.core.example import example_pb2
  from tensorflow.core.example import feature_pb2
  import keras
  from keras import combinations
  from keras import keras_parameterized
  from keras import regularizers
  from keras import testing_utils
  from keras.feature_column.dense_features import DenseFeatures
  from keras.saving.saved_model import load as keras_load
  from keras.saving.saved_model import save_impl as keras_save
  from keras.utils import control_flow_util
  from keras.utils import generic_utils
  from keras.utils import tf_contextlib
  from keras.utils import tf_inspect
  
  
  class LayerWithLearningPhase(keras.engine.base_layer.Layer):
  
    def build(self, input_shape):
      self.input_spec = keras.layers.InputSpec(shape=[None] * len(input_shape))
      self.built = True
  
    def call(self, x, training=None):
      if training is None:
        training = keras.backend.learning_phase()
      output = control_flow_util.smart_cond(training, lambda: x * 0,
                                            lambda: tf.identity(x))
      if not tf.executing_eagerly():
        output._uses_learning_phase = True  # pylint: disable=protected-access
      return output
  
    def compute_output_shape(self, input_shape):
      return input_shape
  
    @property
    def _use_input_spec_as_call_signature(self):
      return True
  
  
  class LayerWithLoss(keras.layers.Layer):
  
    def call(self, inputs):
      self.add_loss(tf.reduce_sum(inputs), inputs=inputs)
      return inputs * 2
  
  
  class LayerWithUpdate(keras.layers.Layer):
  
    def build(self, _):
      self.v = self.add_weight(
          'v',
          shape=[],
          initializer=keras.initializers.zeros,
          trainable=False,
          dtype=tf.float32)
  
    def call(self, inputs, training=True):
      if training:
        self.add_update(self.v.assign_add(1.))
      return inputs * 2.
  
  
  @generic_utils.register_keras_serializable('Testing')
  class GlobalLayerThatShouldFailIfNotAdded(keras.layers.Layer):
    _must_restore_from_config = True
  
  
  @keras_parameterized.run_all_keras_modes
  class TestSavedModelFormatAllModes(keras_parameterized.TestCase):
  
    def _save_model_dir(self, dirname='saved_model'):
      temp_dir = self.get_temp_dir()
      self.addCleanup(shutil.rmtree, temp_dir, ignore_errors=True)
      return os.path.join(temp_dir, dirname)
  
    def _test_save_and_load(self, use_dataset=False):
      model = testing_utils.get_small_mlp(1, 4, input_dim=3)
      model.layers[-1].activity_regularizer = regularizers.get('l2')
      model.activity_regularizer = regularizers.get('l2')
      model.compile(
          loss='mse',
          optimizer='rmsprop')
      def callable_loss():
        return tf.reduce_sum(model.weights[0])
      model.add_loss(callable_loss)
  
      x = np.random.random((1, 3))
      y = np.random.random((1, 4))
  
      if not tf.__internal__.tf2.enabled():
        # The layer autocast behavior only runs when autocast is enabled, so
        # in V1, the numpy inputs still need to be cast to float32.
        x = x.astype(np.float32)
        y = y.astype(np.float32)
  
      if use_dataset:
        dataset = tf.data.Dataset.from_tensor_slices((x, y)).batch(1)
        model.fit(dataset)
      else:
        model.train_on_batch(x, y)
  
      saved_model_dir = self._save_model_dir()
      tf.saved_model.save(model, saved_model_dir)
      loaded = keras_load.load(saved_model_dir)
      self.evaluate(tf.compat.v1.variables_initializer(loaded.variables))
      self.assertAllClose(self.evaluate(model.weights),
                          self.evaluate(loaded.weights))
  
      input_arr = tf.constant(
          np.random.random((1, 3)).astype(np.float32))
      self.assertAllClose(self.evaluate(model(input_arr)),
                          self.evaluate(loaded(input_arr)))
      # Validate losses. The order of conditional losses may change between the
      # model and loaded model, so sort the losses first.
      if tf.executing_eagerly():
        self.assertAllClose(sorted(self.evaluate(model.losses)),
                            sorted(self.evaluate(loaded.losses)))
      else:
        self.assertAllClose(self.evaluate(model.get_losses_for(None)),
                            self.evaluate(loaded.get_losses_for(None)))
        self.assertAllClose(
            sorted(self.evaluate(model.get_losses_for(input_arr))),
            sorted(self.evaluate(loaded.get_losses_for(input_arr))))
  
    @keras_parameterized.run_with_all_model_types
    def test_model_save_and_load(self):
      self._test_save_and_load(use_dataset=True)
  
    @keras_parameterized.run_with_all_model_types
    def test_model_save_and_load_dataset(self):
      self._test_save_and_load(use_dataset=True)
  
    def test_trainable_weights(self):
      layer = keras.layers.Dense(4, name='custom_layer')
      layer.build([3,])
      layer.add_weight(
          'extra_weight', shape=[],
          initializer=tf.compat.v1.constant_initializer(11),
          trainable=True)
      layer.add_weight(
          'extra_weight_2', shape=[],
          initializer=tf.compat.v1.constant_initializer(12),
          trainable=False)
  
      saved_model_dir = self._save_model_dir()
      self.evaluate(tf.compat.v1.variables_initializer(layer.variables))
      tf.saved_model.save(layer, saved_model_dir)
      loaded = keras_load.load(saved_model_dir)
      self.evaluate(tf.compat.v1.variables_initializer(loaded.variables))
  
      equal_attrs = ['name', '_expects_training_arg', 'trainable']
      for attr in equal_attrs:
        self.assertEqual(getattr(layer, attr), getattr(loaded, attr))
  
      all_close = ['weights', 'trainable_weights', 'non_trainable_weights']
      for attr in all_close:
        self.assertAllClose(self.evaluate(getattr(layer, attr)),
                            self.evaluate(getattr(loaded, attr)))
  
    def test_maintains_losses(self):
      """Tests that the layer losses do not change before and after export."""
      model = keras.models.Sequential([LayerWithLoss()])
      model.compile(
          loss='mse',
          optimizer='rmsprop')
      input_arr = np.random.random((1, 3))
      target_arr = np.random.random((1, 3))
  
      # Test that symbolic losses are maintained (train_on_batch saves symbolic
      # losses.)
      model.train_on_batch(input_arr, target_arr)
      previous_losses = model.losses[:]
  
      saved_model_dir = self._save_model_dir()
      tf.saved_model.save(model, saved_model_dir)
  
      with previous_losses[0].graph.as_default():
        # If we try to compare symbolic Tensors in eager mode assertAllEqual will
        # return False even if they are the same Tensor.
        self.assertAllEqual(previous_losses, model.losses)
  
      if tf.executing_eagerly():
        # Test that eager losses are maintained.
        model(input_arr)  # Calls model eagerly, creating eager losses.
        previous_losses = model.losses[:]
        tf.saved_model.save(model, saved_model_dir)
        self.assertAllEqual(previous_losses, model.losses)
  
    def test_layer_with_learning_phase(self):
      layer = LayerWithLearningPhase()
      layer.build([None, None])
      saved_model_dir = self._save_model_dir()
      tf.saved_model.save(layer, saved_model_dir)
      loaded = keras_load.load(saved_model_dir)
      input_arr = tf.ones((4, 3))
  
      # Run the layer, and use the keras backend learning phase
      keras.backend.set_learning_phase(0)
      self.assertAllEqual(input_arr, loaded(input_arr))
      keras.backend.set_learning_phase(1)
      self.assertAllEqual(tf.zeros((4, 3)), loaded(input_arr))
  
      # Run the layer while explicitly setting the training argument
      self.assertAllEqual(
          input_arr, loaded(input_arr, training=tf.constant(False)))
      self.assertAllEqual(
          tf.zeros((4, 3)),
          loaded(input_arr, training=tf.constant(True)))
  
    @keras_parameterized.run_with_all_model_types
    def test_standard_loader(self):
      model = testing_utils.get_small_mlp(1, 4, input_dim=3)
      model.activity_regularizer = regularizers.get('l2')
      def eager_loss():
        return tf.reduce_sum(model.weights[0])
      model.add_loss(eager_loss)
  
      # Call predict to ensure that all layers are built and inputs are set.
      model.predict(np.random.random((1, 3)).astype(np.float32))
      saved_model_dir = self._save_model_dir()
  
      tf.saved_model.save(model, saved_model_dir)
  
      loaded = tf.saved_model.load(saved_model_dir)
      self.evaluate(tf.compat.v1.variables_initializer(loaded.variables))
      all_close = ['variables', 'trainable_variables',
                   'non_trainable_variables']
      for attr in all_close:
        self.assertAllClose(self.evaluate(getattr(model, attr)),
                            self.evaluate(getattr(loaded.keras_api, attr)))
      self.assertLen(loaded.regularization_losses, 1)
      expected_layers = len(model.layers)
      self.assertEqual(expected_layers, len(loaded.keras_api.layers))
      input_arr = tf.ones((4, 3))
      self.assertAllClose(self.evaluate(model(input_arr)),
                          self.evaluate(loaded(input_arr, training=False)))
  
    @keras_parameterized.run_with_all_model_types
    def test_compiled_model(self):
      # TODO(b/134519980): Issue with model.fit if the model call function uses
      # a tf.function (Graph mode only).
      if not tf.executing_eagerly():
        return
  
      input_arr = np.random.random((1, 3))
      target_arr = np.random.random((1, 4))
  
      model = testing_utils.get_small_mlp(1, 4, input_dim=3)
      expected_predict = model.predict(input_arr)
  
      # Compile and save model.
      model.compile('rmsprop', 'mse')
      saved_model_dir = self._save_model_dir()
      tf.saved_model.save(model, saved_model_dir)
  
      loaded = keras_load.load(saved_model_dir)
      actual_predict = loaded.predict(input_arr)
      self.assertAllClose(expected_predict, actual_predict)
  
      loss_before = loaded.evaluate(input_arr, target_arr)
      loaded.fit(input_arr, target_arr)
      loss_after = loaded.evaluate(input_arr, target_arr)
      self.assertLess(loss_after, loss_before)
      predict = loaded.predict(input_arr)
  
      ckpt_path = os.path.join(self.get_temp_dir(), 'weights')
      loaded.save_weights(ckpt_path)
  
      # Ensure that the checkpoint is compatible with the original model.
      model.load_weights(ckpt_path)
      self.assertAllClose(predict, model.predict(input_arr))
  
    def test_metadata_input_spec(self):
      class LayerWithNestedSpec(keras.layers.Layer):
  
        def __init__(self):
          super(LayerWithNestedSpec, self).__init__()
          self.input_spec = {
              'a': keras.layers.InputSpec(max_ndim=3, axes={-1: 2}),
              'b': keras.layers.InputSpec(shape=(None, 2, 3), dtype='float16')}
  
        @property
        def _use_input_spec_as_call_signature(self):
          return True
  
      layer = LayerWithNestedSpec()
      saved_model_dir = self._save_model_dir()
      tf.saved_model.save(layer, saved_model_dir)
      loaded = keras_load.load(saved_model_dir)
      self.assertEqual(3, loaded.input_spec['a'].max_ndim)
      self.assertEqual({-1: 2}, loaded.input_spec['a'].axes)
      self.assertAllEqual([None, 2, 3], loaded.input_spec['b'].shape)
      self.assertEqual('float16', loaded.input_spec['b'].dtype)
  
    def test_must_restore_from_config_fails_if_layer_is_not_in_scope(self):
  
      class LayerThatShouldFailIfNotAdded(keras.layers.Layer):
        _must_restore_from_config = True
  
      layer = LayerThatShouldFailIfNotAdded()
      saved_model_dir = self._save_model_dir()
      tf.saved_model.save(layer, saved_model_dir)
      with self.assertRaisesRegex(RuntimeError, 'Unable to restore a layer of'):
        _ = keras_load.load(saved_model_dir)
  
    def test_must_restore_from_config_custom_object_scope(self):
  
      class LayerThatShouldFailIfNotAdded(keras.layers.Layer):
        _must_restore_from_config = True
  
      layer = LayerThatShouldFailIfNotAdded()
      saved_model_dir = self._save_model_dir()
      tf.saved_model.save(layer, saved_model_dir)
      with generic_utils.CustomObjectScope(
          {'LayerThatShouldFailIfNotAdded': LayerThatShouldFailIfNotAdded}):
        _ = keras_load.load(saved_model_dir)
  
    def test_must_restore_from_config_registration(self):
      layer = GlobalLayerThatShouldFailIfNotAdded()
      saved_model_dir = self._save_model_dir()
      tf.saved_model.save(layer, saved_model_dir)
      _ = keras_load.load(saved_model_dir)
  
    def test_multi_input_model(self):
      input_1 = keras.layers.Input(shape=(3,))
      input_2 = keras.layers.Input(shape=(5,))
      model = keras.Model([input_1, input_2], [input_1, input_2])
      saved_model_dir = self._save_model_dir()
  
      model.save(saved_model_dir, save_format='tf')
      loaded = keras_load.load(saved_model_dir)
      input_arr_1 = np.random.random((1, 3)).astype('float32')
      input_arr_2 = np.random.random((1, 5)).astype('float32')
  
      outputs = loaded([input_arr_1, input_arr_2])
      self.assertAllEqual(input_arr_1, outputs[0])
      self.assertAllEqual(input_arr_2, outputs[1])
  
    def test_revived_sequential(self):
      model = keras.models.Sequential()
      model.add(keras.layers.Dense(5, input_shape=(3,),
                                   kernel_regularizer=regularizers.get('l2')))
      model.add(keras.layers.Dense(2, kernel_regularizer=regularizers.get('l2')))
  
      self.evaluate(tf.compat.v1.variables_initializer(model.variables))
  
      saved_model_dir = self._save_model_dir()
      model.save(saved_model_dir, save_format='tf')
      loaded = keras_load.load(saved_model_dir)
  
      self.assertLen(loaded.layers, 2)
      self.assertLen(loaded.losses, 2)
  
      loaded.pop()
  
      self.assertLen(loaded.layers, 1)
      self.assertLen(loaded.losses, 1)
  
      loaded.add(keras.layers.Dense(2, kernel_regularizer=regularizers.get('l2')))
  
      self.assertLen(loaded.layers, 2)
      self.assertLen(loaded.losses, 2)
  
    def testBatchNormUpdates(self):
      model = keras.models.Sequential(
          keras.layers.BatchNormalization(input_shape=(1,)))
      self.evaluate(tf.compat.v1.variables_initializer(model.variables))
      saved_model_dir = self._save_model_dir()
  
      # TODO(kathywu): Re-enable this check after removing the tf.saved_model.save
      # metadata warning.
      # with self.captureWritesToStream(sys.stderr) as captured_logs:
      model.save(saved_model_dir, save_format='tf')
      loaded = keras_load.load(saved_model_dir)
  
      # Assert that saving does not log deprecation warnings
      # (even if it needs to set learning phase for compat reasons)
      # if context.executing_eagerly():
      #   self.assertNotIn('deprecated', captured_logs.contents())
  
      input_arr = tf.constant([[11], [12], [13]], dtype=tf.float32)
      input_arr2 = tf.constant([[14], [15], [16]], dtype=tf.float32)
      self.assertAllClose(self.evaluate(loaded.layers[-1].moving_mean), [0])
  
      self.evaluate(loaded(input_arr, training=True))
      if not tf.executing_eagerly():
        self.evaluate(loaded.get_updates_for(input_arr))
      self.assertAllClose(self.evaluate(loaded.layers[-1].moving_mean), [0.12])
  
      self.evaluate(loaded(input_arr2, training=False))
      if not tf.executing_eagerly():
        self.evaluate(loaded.get_updates_for(input_arr2))
      self.assertAllClose(self.evaluate(loaded.layers[-1].moving_mean), [0.12])
  
    def testDisablingBatchNormTrainableBeforeSaving(self):
      # We disable trainable on the batchnorm layers before saving
      model = keras.models.Sequential(
          keras.layers.BatchNormalization(input_shape=(1,)))
      model.trainable = False
      self.evaluate(tf.compat.v1.variables_initializer(model.variables))
      saved_model_dir = self._save_model_dir()
      model.save(saved_model_dir, save_format='tf')
      loaded = keras_load.load(saved_model_dir)
      self.evaluate(tf.compat.v1.variables_initializer(loaded.variables))
      input_arr = tf.constant([[11], [12], [13]], dtype=tf.float32)
      input_arr2 = tf.constant([[14], [15], [16]], dtype=tf.float32)
      self.assertAllClose(self.evaluate(loaded.layers[-1].moving_mean), [0])
  
      # Trainable should still be disabled after loading
      self.evaluate(loaded(input_arr, training=True))
      if not tf.executing_eagerly():
        self.evaluate(loaded.get_updates_for(input_arr))
      self.assertAllClose(self.evaluate(loaded.layers[-1].moving_mean), [0.0])
  
      # Re-enabling trainable on the loaded model should cause the batchnorm
      # layer to start training again.
      # Note: this only works in v2.
      if tf.executing_eagerly():
        loaded.trainable = True
        self.evaluate(loaded(input_arr, training=True))
        self.assertAllClose(self.evaluate(loaded.layers[-1].moving_mean), [0.12])
  
        self.evaluate(loaded(input_arr2, training=False))
        self.assertAllClose(self.evaluate(loaded.layers[-1].moving_mean), [0.12])
  
    def testSaveWithSignatures(self):
      model = keras.models.Sequential()
      model.add(keras.layers.Dense(5, input_shape=(3,),
                                   kernel_regularizer=regularizers.get('l2')))
      model.add(keras.layers.Dropout(0.5))
      model.add(keras.layers.Dense(4, kernel_regularizer=regularizers.get('l2')))
  
      input_arr = np.random.random((2, 3))
      target_arr = np.random.random((2, 4))
  
      model.compile(
          loss='mse',
          optimizer='rmsprop')
      model.train_on_batch(input_arr, target_arr)
  
      @tf.function(input_signature=[tf.TensorSpec((None, 3))])
      def predict(inputs):
        return {'predictions': model(inputs)}
  
      feature_configs = {
          'inputs': tf.io.FixedLenFeature(
              shape=[2, 3], dtype=tf.float32)}
  
      @tf.function(
          input_signature=[tf.TensorSpec([None], tf.string)])
      def parse_and_predict(examples):
        features = tf.compat.v1.parse_single_example(examples[0], feature_configs)
        return {'predictions': model(features['inputs']),
                'layer_1_outputs': model.layers[0](features['inputs'])}
  
      saved_model_dir = self._save_model_dir()
      model.save(saved_model_dir, save_format='tf', signatures={
          'predict': predict,
          'parse_and_predict': parse_and_predict})
      model.save('/tmp/saved', save_format='tf', signatures={
          'predict': predict,
          'parse_and_predict': parse_and_predict})
  
      loaded = keras_load.load(saved_model_dir)
  
      self.assertAllClose(
          model.predict(input_arr),
          loaded.signatures['predict'](tf.convert_to_tensor(
              input_arr.astype('float32')))['predictions'])
  
      feature = {
          'inputs': feature_pb2.Feature(
              float_list=feature_pb2.FloatList(
                  value=input_arr.astype('float32').flatten()))}
      example = example_pb2.Example(
          features=feature_pb2.Features(feature=feature))
      outputs = loaded.signatures['parse_and_predict'](
          tf.convert_to_tensor([example.SerializeToString()]))
      self.assertAllClose(model.predict(input_arr), outputs['predictions'])
      self.assertAllClose(model.layers[0](input_arr), outputs['layer_1_outputs'])
  
    def testTrainingDefaults(self):
      def assert_training_default(fn, default_value):
        arg_spec = tf_inspect.getfullargspec(fn)
        index = len(arg_spec.args) - arg_spec.args.index('training')
        self.assertEqual(arg_spec.defaults[-index], default_value)
  
      class LayerWithTrainingRequiredArg(keras.engine.base_layer.Layer):
  
        def call(self, inputs, training):
          return control_flow_util.smart_cond(training, lambda: inputs * 0,
                                              lambda: tf.identity(inputs))
  
      class LayerWithTrainingDefaultTrue(keras.engine.base_layer.Layer):
  
        def call(self, inputs, training=True):
          return control_flow_util.smart_cond(training, lambda: inputs * 0,
                                              lambda: tf.identity(inputs))
  
      class Model(keras.models.Model):
  
        def __init__(self):
          super(Model, self).__init__()
          self.layer_with_training_default_none = LayerWithLearningPhase()
          self.layer_with_training_default_true = LayerWithTrainingDefaultTrue()
          self.layer_with_required_training_arg = LayerWithTrainingRequiredArg()
  
        def call(self, inputs):
          x = self.layer_with_training_default_none(inputs)
          x += self.layer_with_training_default_true(inputs)
          x += self.layer_with_required_training_arg(inputs, False)
          return x
  
      model = Model()
      # Build and set model inputs
      model.predict(np.ones([1, 3]).astype('float32'))
      saved_model_dir = self._save_model_dir()
      model.save(saved_model_dir, save_format='tf')
      load = tf.saved_model.load(saved_model_dir)
  
      # Ensure that the Keras loader is able to load and build the model.
      _ = keras_load.load(saved_model_dir)
  
      assert_training_default(load.__call__, False)
      assert_training_default(
          load.layer_with_training_default_none.__call__, False)
      assert_training_default(
          load.layer_with_training_default_true.__call__, True)
  
      # Assert that there are no defaults for layer with required training arg
      arg_spec = tf_inspect.getfullargspec(
          load.layer_with_required_training_arg.__call__)
      self.assertFalse(arg_spec.defaults)  # defaults is None or empty
  
    def testTraceModelWithKwarg(self):
      class Model(keras.models.Model):
  
        def call(self, inputs, keyword=None):
          return tf.identity(inputs)
  
      model = Model()
      prediction = model.predict(np.ones([1, 3]).astype('float32'))
      saved_model_dir = self._save_model_dir()
      model.save(saved_model_dir, save_format='tf')
  
      loaded = keras_load.load(saved_model_dir)
      self.assertAllClose(prediction,
                          loaded.predict(np.ones([1, 3]).astype('float32')))
  
    def testFeatureColumns(self):
      # TODO(b/120099662): Error with table initialization with Keras models in
      # graph mode.
      if tf.executing_eagerly():
        numeric = tf.feature_column.numeric_column('a')
        bucketized = tf.feature_column.bucketized_column(numeric, boundaries=[5, 10, 15])
        cat_vocab = tf.feature_column.categorical_column_with_vocabulary_list(
            'b', ['1', '2', '3'])
        one_hot = tf.feature_column.indicator_column(cat_vocab)
        embedding = tf.feature_column.embedding_column(cat_vocab, dimension=8)
        feature_layer = DenseFeatures([bucketized, one_hot, embedding])
        model = keras.models.Sequential(feature_layer)
  
        features = {'a': np.array([13, 15]), 'b': np.array(['1', '2'])}
        predictions = model.predict(features)
  
        saved_model_dir = self._save_model_dir()
        model.save(saved_model_dir, save_format='tf')
        loaded = keras_load.load(saved_model_dir)
        loaded_predictions = loaded.predict(features)
        self.assertAllClose(predictions, loaded_predictions)
  
    def testSaveTensorKwarg(self):
  
      class LayerWithTensorKwarg(keras.layers.Layer):
  
        def call(self, inputs, tensor=None):
          if tensor is not None:
            return inputs * tf.cast(tensor, tf.float32)
          else:
            return inputs
  
      t = self.evaluate(tf.sequence_mask(1))
      inputs = keras.layers.Input(shape=(3))
      model = keras.models.Model(inputs, LayerWithTensorKwarg()(inputs, t))
  
      input_arr = np.random.random((1, 3))
      predictions = model.predict(input_arr)
  
      saved_model_dir = self._save_model_dir()
      model.save(saved_model_dir, save_format='tf')
      loaded = keras_load.load(saved_model_dir)
      loaded_predictions = loaded.predict(input_arr)
      self.assertAllClose(predictions, loaded_predictions)
  
    def testModelWithTfFunctionCall(self):
      class Subclass(keras.models.Model):
  
        @tf.function
        def call(self, inputs, training=False):
          return inputs * tf.cast(training, tf.float32)
  
      model = Subclass()
      model.predict(tf.ones((1, 2)), steps=1)
      saved_model_dir = self._save_model_dir()
      model.save(saved_model_dir, save_format='tf')
      loaded = keras_load.load(saved_model_dir)
      self.assertAllEqual(
          [[1, 5]],
          self.evaluate(loaded(tf.constant([[1, 5.]]), training=True)))
      self.assertAllEqual(
          [[0, 0]],
          self.evaluate(loaded(tf.constant([[1, 5.]]), training=False)))
  
    def testReviveFunctionalModel(self):
  
      class CustomAdd(keras.layers.Add):
  
        def build(self, input_shape):
          self.w = self.add_weight('w', shape=[])
          super(CustomAdd, self).build(input_shape)
  
        def call(self, inputs):
          outputs = super(CustomAdd, self).call(inputs)
          return outputs * self.w
  
      input1 = keras.layers.Input(shape=(None, 3), name='input_1')
      input2 = keras.layers.Input(shape=(None, 3), name='input_2')
  
      d = keras.layers.Dense(4, name='dense_with_two_inbound_nodes')
      output1 = d(input1)
      output2 = d(input2)
  
      # Use a custom layer in this model to ensure that layers aren't being
      # recreated directly from the config.
      outputs = CustomAdd(name='custom')([output1, output2])
      model = keras.models.Model([input1, input2], outputs, name='save_model')
  
      self.evaluate(tf.compat.v1.variables_initializer(model.variables))
      saved_model_dir = self._save_model_dir()
      model.save(saved_model_dir, save_format='tf')
  
      loaded = keras_load.load(saved_model_dir)
      self.assertEqual('save_model', loaded.name)
      self.assertLen(
          loaded.get_layer('dense_with_two_inbound_nodes')._inbound_nodes, 2)
      self.assertEqual('CustomAdd', type(loaded.get_layer('custom')).__name__)
      self.assertLen(loaded.get_layer('custom').weights, 1)
  
    def _testAddUpdate(self, scope):
      with scope:
        layer_with_update = LayerWithUpdate()
        model = testing_utils.get_model_from_layers([layer_with_update],
                                                    input_shape=(3,))
  
        x = np.ones((10, 3))
        if testing_utils.get_model_type() == 'subclass':
          model.predict(x, batch_size=10)
        self.evaluate(tf.compat.v1.variables_initializer(model.variables))
        saved_model_dir = self._save_model_dir()
        model.save(saved_model_dir, save_format='tf')
  
      loaded = keras_load.load(saved_model_dir)
      loaded_layer = loaded.layers[-1]
      self.evaluate(tf.compat.v1.variables_initializer(loaded.variables))
      self.assertEqual(self.evaluate(loaded_layer.v), 0.)
  
      loaded.compile('sgd', 'mse')
      loaded.fit(x, x, batch_size=10)
      self.assertEqual(self.evaluate(loaded_layer.v), 1.)
  
    @keras_parameterized.run_with_all_model_types
    def testSaveLayerWithUpdates(self):
      @tf_contextlib.contextmanager
      def nullcontextmanager():
        yield
      self._testAddUpdate(nullcontextmanager())
  
    @keras_parameterized.run_with_all_model_types
    def testSaveInStrategyScope(self):
      self._testAddUpdate(tf.distribute.MirroredStrategy().scope())
  
    def testSaveTimeDistributedLayer(self):
      model = keras.Sequential([
          keras.layers.TimeDistributed(
              keras.layers.Dense(1, kernel_regularizer=regularizers.get('l2')),
              input_shape=(None, 1))])
      predictions = model.predict_on_batch(tf.ones((3, 2, 1)))
  
      saved_model_dir = self._save_model_dir()
      model.save(saved_model_dir, save_format='tf')
  
      loaded = keras_load.load(saved_model_dir)
      self.assertAllClose(loaded.predict_on_batch(tf.ones((3, 2, 1))),
                          predictions)
  
    @parameterized.named_parameters([
        ('with_unrolling', True),
        ('no_unrolling', False)
    ])
    def testSaveStatefulRNN(self, unroll):
      batch = 12
      timesteps = 10
      input_dim = 8
      input_arr = np.ones((batch, timesteps, input_dim)).astype('float32')
  
      cells = [keras.layers.LSTMCell(32), keras.layers.LSTMCell(64)]
      if unroll:
        x = keras.Input(batch_shape=(batch, timesteps, input_dim))
      else:
        x = keras.Input(batch_shape=(batch, None, input_dim))
      layer = keras.layers.RNN(cells, stateful=True, unroll=unroll)
      y = layer(x)
  
      model = keras.Model(x, y)
      model.compile('rmsprop', 'mse',
                    run_eagerly=testing_utils.should_run_eagerly())
      model.train_on_batch(
          np.zeros((batch, timesteps, input_dim)).astype('float32'),
          np.zeros((batch, 64)).astype('float32'))
  
      saved_model_dir = self._save_model_dir()
      tf.saved_model.save(model, saved_model_dir)
  
      loaded = keras_load.load(saved_model_dir)
      loaded_layer = loaded.layers[1]
  
      if not tf.executing_eagerly():
        keras.backend.get_session()  # force variable initialization
  
      self.assertAllClose(layer.states, loaded_layer.states)
      self.assertAllClose(model(input_arr), loaded(input_arr))
  
    def testSaveStatelessConvLSTM2D(self):
      data_format = 'channels_first'
      batch, timesteps, channels, rows, cols = 12, 10, 8, 4, 4
      input_arr = np.ones(
          (batch, timesteps, channels, rows, cols)).astype('float32')
      layer = keras.layers.ConvLSTM2D(
          filters=16, kernel_size=(1, 1), data_format=data_format)
      x = keras.Input(batch_shape=(batch, timesteps, channels, rows, cols))
      y = layer(x)
      model = keras.Model(x, y)
  
      predict_1 = model(input_arr)
      saved_model_dir = self._save_model_dir()
      tf.saved_model.save(model, saved_model_dir)
      del model
  
      loaded = keras_load.load(saved_model_dir)
      predict_2 = loaded(input_arr)
      self.assertAllClose(predict_1, predict_2)
  
    def testSaveWithRaggedInputs(self):
  
      class EmbeddingMerger(keras.layers.Layer):
  
        def __init__(self, list_features, **kwargs):
          super().__init__(**kwargs)
          self._supports_ragged_inputs = True
          self.embeddings = {
              feature: keras.layers.Embedding(10, 3) for feature in list_features}
          self.mean = keras.layers.Lambda(
              tf.reduce_mean, arguments=dict(axis=1))
  
        def call(self, inputs):
          tensors = [self.embeddings[col](inputs[col]) for col in inputs]
          tensors = [self.mean(inp) for inp in tensors]
          return keras.layers.Add()(tensors)
  
      list_features = ['feature_1', 'feature_2']
      feature_1 = tf.ragged.constant([[0.], [1, 3]])
      feature_2 = tf.ragged.constant([[1., 2], [4]])
      f = {'feature_1': feature_1,
           'feature_2': feature_2}
      f_inputs = {
          'feature_1': keras.Input(shape=(None,), name='feature_1', ragged=True),
          'feature_2': keras.Input(shape=(None,), name='feature_2', ragged=True)}
  
      out = EmbeddingMerger(list_features)(f_inputs)
      model = keras.Model(f_inputs, out)
      self.evaluate(tf.compat.v1.variables_initializer(model.variables))
      saved_model_dir = self._save_model_dir()
      tf.saved_model.save(model, saved_model_dir)
  
      loaded = keras_load.load(saved_model_dir)
      self.evaluate(tf.compat.v1.variables_initializer(loaded.variables))
      self.assertAllClose(model.predict(f), loaded.predict(f))
  
  
  class TestSavedModelFormat(tf.test.TestCase):
  
    def _save_model_dir(self, dirname='saved_model'):
      temp_dir = self.get_temp_dir()
      self.addCleanup(shutil.rmtree, temp_dir, ignore_errors=True)
      return os.path.join(temp_dir, dirname)
  
    def test_load_with_partially_failed_serialization(self):
  
      class BadCustomLayer(keras.layers.Layer):
  
        def __call__(self, inputs):
          return inputs
  
      class Model(keras.models.Model):
  
        def __init__(self):
          super(Model, self).__init__()
          self.layer = BadCustomLayer()
  
        @tf.function(
            input_signature=[tf.TensorSpec([None, 1])])
        def call(self, inputs):
          return self.layer(inputs)
  
      model = Model()
      inp = tf.constant([[1.0]])
      model(inp)
      saved_model_dir = self._save_model_dir()
      tf.saved_model.save(model, saved_model_dir)
  
      loaded = keras_load.load(saved_model_dir)
      self.assertAllEqual([[1.0]], self.evaluate(loaded(inp)))
      with self.assertRaisesRegex(ValueError, 'call function was not serialized'):
        loaded.layer(inp)
  
    def test_save_without_tracing(self):
  
      class DoNotTrace(keras.layers.Layer):
  
        def __init__(self):
          super(DoNotTrace, self).__init__()
          self.input_spec = keras.layers.InputSpec(shape=[None])
          self.built = True
  
        def call(self, inputs):
          raise ValueError('I said do not trace')
  
        def get_config(self):
          return {}
  
        @property
        def _use_input_spec_as_call_signature(self):
          return True
  
      root = keras.models.Sequential()
      root.add(keras.layers.Input(shape=(3,)))
      root.attached_layer = DoNotTrace()
  
      saved_model_dir = self._save_model_dir()
  
      # With the default settings, the call function is traced.
      with self.assertRaisesRegex(ValueError, 'do not trace'):
        root.save(saved_model_dir, save_format='tf')
  
      # When saving the config only, the layer call function should not be not
      # traced.
      root.save(saved_model_dir, save_format='tf', save_traces=False)
      loaded = tf.saved_model.load(saved_model_dir)
      self.assertTrue(hasattr(loaded, 'attached_layer'))
  
      # This should raise an error when loaded without the custom object
      loaded = keras_load.load(saved_model_dir)
      with self.assertRaisesRegex(ValueError, 'Cannot call custom layer'):
        loaded.attached_layer(tf.constant([1.]))
  
      # Try loading with the custom objects
      with generic_utils.CustomObjectScope({'DoNotTrace': DoNotTrace}):
        loaded = keras_load.load(saved_model_dir)
      with self.assertRaisesRegex(ValueError, 'I said do not trace'):
        loaded.attached_layer(tf.constant([1.]))
  
  
  class TestLayerCallTracing(tf.test.TestCase, parameterized.TestCase):
  
    def test_functions_have_same_trace(self):
  
      class Layer(keras.engine.base_layer.Layer):
  
        def call(self, inputs):
          return inputs
  
        def call2(self, inputs):
          return inputs * 2
  
      layer = Layer()
      call_collection = keras_save.LayerCallCollection(layer)
      fn = call_collection.add_function(layer.call, 'call')
      fn2 = call_collection.add_function(layer.call2, 'call2')
  
      fn(np.ones((2, 3)))
      fn(np.ones((4, 5)))
  
      self.assertLen(fn._list_all_concrete_functions_for_serialization(), 2)
      self.assertLen(fn2._list_all_concrete_functions_for_serialization(), 2)
  
      # Check that the shapes are correct
      self.assertEqual(
          {(2, 3), (4, 5)},
          set(tuple(c.structured_input_signature[0][0].shape.as_list())
              for c in fn2._list_all_concrete_functions_for_serialization()))
  
    def test_training_arg_replacement(self):
  
      def assert_num_traces(layer_cls, training_keyword):
        layer = layer_cls()
        call_collection = keras_save.LayerCallCollection(layer)
        fn = call_collection.add_function(layer.call, 'call')
  
        fn(np.ones((2, 3)), training=True)
        self.assertLen(fn._list_all_concrete_functions_for_serialization(), 2)
  
        fn(np.ones((2, 4)), training=False)
        self.assertLen(fn._list_all_concrete_functions_for_serialization(), 4)
  
        if training_keyword:
          fn(np.ones((2, 5)), True)
          self.assertLen(fn._list_all_concrete_functions_for_serialization(), 6)
          fn(np.ones((2, 6)))
          self.assertLen(fn._list_all_concrete_functions_for_serialization(), 8)
  
      class LayerWithTrainingKeyword(keras.engine.base_layer.Layer):
  
        def call(self, inputs, training=False):
          return inputs * training
  
      assert_num_traces(LayerWithTrainingKeyword, training_keyword=True)
  
      class LayerWithKwargs(keras.engine.base_layer.Layer):
  
        def call(self, inputs, **kwargs):
          return inputs * kwargs['training']
  
      assert_num_traces(LayerWithKwargs, training_keyword=False)
  
      class LayerWithChildLayer(keras.engine.base_layer.Layer):
  
        def __init__(self):
          self.child = LayerWithKwargs()
          super(LayerWithChildLayer, self).__init__()
  
        def call(self, inputs):
          return self.child(inputs)
  
      assert_num_traces(LayerWithChildLayer, training_keyword=False)
  
    @combinations.generate(combinations.combine(mode=['graph', 'eager']))
    def test_maintains_losses(self):
      layer = LayerWithLoss()
      layer(np.ones((2, 3)))
      previous_losses = layer.losses[:]
  
      call_collection = keras_save.LayerCallCollection(layer)
      fn = call_collection.add_function(layer.call, 'call')
      fn(np.ones((2, 3)))
  
      self.assertAllEqual(previous_losses, layer.losses)
  
  
  @combinations.generate(combinations.combine(mode=['graph', 'eager']))
  class MetricTest(tf.test.TestCase, parameterized.TestCase):
  
    def _save_model_dir(self, dirname='saved_model'):
      temp_dir = self.get_temp_dir()
      self.addCleanup(shutil.rmtree, temp_dir, ignore_errors=True)
      return os.path.join(temp_dir, dirname)
  
    def generate_inputs(self, num_tensor_args, shape=(1, 5)):
      return [
          np.random.uniform(0, 1, shape).astype('float32')
          for _ in range(num_tensor_args)
      ]
  
    def _test_metric_save_and_load(self,
                                   metric,
                                   save_dir,
                                   num_tensor_args,
                                   shape=(1, 5),
                                   test_sample_weight=True):
      with self.cached_session():
        tf.saved_model.save(metric, save_dir)
        loaded = keras_load.load(save_dir)
        self.evaluate([v.initializer for v in loaded.variables])
        self.assertEqual(metric.name, loaded.name)
        self.assertEqual(metric.dtype, loaded.dtype)
  
        inputs = self.generate_inputs(num_tensor_args, shape)
        actual = self.evaluate(metric(*inputs))
        self.assertAllClose(actual, loaded(*inputs))
        self.assertAllClose(metric.variables, loaded.variables)
  
        # Test with separate calls to update state and result.
        inputs = self.generate_inputs(num_tensor_args, shape)
        self.evaluate(metric.update_state(*inputs))
        self.evaluate(loaded.update_state(*inputs))
        actual = self.evaluate(metric.result())
        self.assertAllClose(actual, loaded.result())
  
        if test_sample_weight:
          # Test with sample weights input.
          inputs = self.generate_inputs(num_tensor_args, shape)
          sample_weight = self.generate_inputs(1, [])[0]
          inputs.append(sample_weight)
  
          actual = self.evaluate(metric(*inputs))
          self.assertAllClose(actual, loaded(*inputs))
        return loaded
  
    @parameterized.named_parameters([
        ('mean', keras.metrics.Mean, 1, (1, 5)),
        ('false_positives', keras.metrics.FalsePositives, 2, (1, 5)),
        ('precision_at_top_k', keras.metrics.Precision, 2, (2, 3, 4), {
            'top_k': 2,
            'class_id': 1
        }),
        ('precision_at_recall', keras.metrics.PrecisionAtRecall, 2, (1, 5), {
            'recall': .8
        }), ('auc', keras.metrics.AUC, 2, (1, 5), {
            'multi_label': True
        }), ('cosine_similarity', keras.metrics.CosineSimilarity, 2, (2, 3, 1))
    ])
    def test_metric(self, metric_cls, num_tensor_args, shape, init_kwargs=None):
      init_kwargs = init_kwargs or {}
      metric = metric_cls(**init_kwargs)
      metric(*self.generate_inputs(num_tensor_args, shape))
      self.evaluate([v.initializer for v in metric.variables])
      loaded = self._test_metric_save_and_load(metric, self._save_model_dir(),
                                               num_tensor_args, shape)
      self.assertEqual(type(loaded), type(metric))
  
    @parameterized.named_parameters([
        ('mean', keras.metrics.Mean, 1, False),
        ('auc', keras.metrics.AUC, 2, False),
        ('mean_tensor', keras.metrics.MeanTensor, 1, True)])
    def test_custom_metric(self, base_cls, num_tensor_args, requires_build):
  
      class CustomMetric(base_cls):
  
        def update_state(self, *args):  # pylint: disable=useless-super-delegation
          # Sometimes built-in metrics return an op in update_state. Custom
          # metrics don't support returning ops, so wrap the update_state method
          # while returning nothing.
          super(CustomMetric, self).update_state(*args)
  
      with self.cached_session():
        metric = CustomMetric()
        save_dir = self._save_model_dir('first_save')
  
        if requires_build:
          metric(*self.generate_inputs(num_tensor_args))  # pylint: disable=not-callable
  
        self.evaluate([v.initializer for v in metric.variables])
  
        with self.assertRaisesRegex(ValueError,
                                    'Unable to restore custom object'):
          self._test_metric_save_and_load(metric, save_dir, num_tensor_args)
        with generic_utils.CustomObjectScope({'CustomMetric': CustomMetric}):
          loaded = self._test_metric_save_and_load(
              metric,
              save_dir,
              num_tensor_args,
              test_sample_weight=False)
  
          self._test_metric_save_and_load(
              loaded,
              self._save_model_dir('second_save'),
              num_tensor_args,
              test_sample_weight=False)
  
    def test_registered_custom_metric(self):
  
      @generic_utils.register_keras_serializable('Testing')
      class CustomMeanMetric(keras.metrics.Mean):
  
        def update_state(self, *args):  # pylint: disable=useless-super-delegation
          # Sometimes built-in metrics return an op in update_state. Custom
          # metrics don't support returning ops, so wrap the update_state method
          # while returning nothing.
          super(CustomMeanMetric, self).update_state(*args)
  
      with self.cached_session():
        metric = CustomMeanMetric()
        save_dir = self._save_model_dir('first_save')
        self.evaluate([v.initializer for v in metric.variables])
        loaded = self._test_metric_save_and_load(
            metric,
            save_dir,
            num_tensor_args=1,
            test_sample_weight=False)
  
        self._test_metric_save_and_load(
            loaded,
            self._save_model_dir('second_save'),
            num_tensor_args=1,
            test_sample_weight=False)
  
    def test_custom_metric_wrapped_call(self):
  
      class NegativeMean(keras.metrics.Mean):
  
        @tf.function(
            input_signature=[tf.TensorSpec(None, tf.float32)])
        def update_state(self, value):
          super(NegativeMean, self).update_state(-value)
  
      metric = NegativeMean()
      self.evaluate([v.initializer for v in metric.variables])
      with generic_utils.CustomObjectScope({'NegativeMean': NegativeMean}):
        self._test_metric_save_and_load(
            metric, self._save_model_dir(), 1, test_sample_weight=False)
  
    @keras_parameterized.run_with_all_model_types
    def test_custom_metric_model(self):
  
      class CustomMetric(keras.metrics.MeanSquaredError):
        pass
  
      model = testing_utils.get_small_mlp(1, 4, input_dim=3)
      model.compile(loss='mse', optimizer='rmsprop', metrics=[CustomMetric()])
  
      saved_model_dir = self._save_model_dir()
      tf.saved_model.save(model, saved_model_dir)
      with self.assertRaisesRegex(ValueError, 'custom_objects'):
        keras_load.load(saved_model_dir)
  
      keras_load.load(saved_model_dir, compile=False)
  
  
  if __name__ == '__main__':
    tf.test.main()
  