text: # Copyright 2018 Google LLC
  #
  # 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.
  # ==============================================================================
  """Unit tests for artifact conversion to and from Python keras."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import glob
  import json
  import os
  import shutil
  import tempfile
  import unittest
  
  import h5py
  import numpy as np
  import tensorflow.compat.v2 as tf
  
  from tensorflowjs import version
  from tensorflowjs.converters import keras_h5_conversion as conversion
  
  
  class ConvertH5WeightsTest(unittest.TestCase):
  
    def setUp(self):
      self._tmp_dir = tempfile.mkdtemp()
      super(ConvertH5WeightsTest, self).setUp()
  
    def tearDown(self):
      if os.path.isdir(self._tmp_dir):
        shutil.rmtree(self._tmp_dir)
      super(ConvertH5WeightsTest, self).tearDown()
  
    def testConvertWeightsFromSimpleModelNoSplitByLayer(self):
      input_tensor = tf.keras.layers.Input((3,))
      dense1 = tf.keras.layers.Dense(
          4, use_bias=True, kernel_initializer='ones', bias_initializer='zeros',
          name='MyDense10')(input_tensor)
      output = tf.keras.layers.Dense(
          2, use_bias=False, kernel_initializer='ones', name='MyDense20')(dense1)
      model = tf.keras.models.Model(inputs=[input_tensor], outputs=[output])
      h5_path = os.path.join(self._tmp_dir, 'MyModel.h5')
      model.save_weights(h5_path)
  
      # Load the saved weights as a JSON string.
      groups = conversion.h5_weights_to_tfjs_format(h5py.File(h5_path))
  
      # Check the loaded weights.
      # Due to the default `split_by_layer=True`, there should be only one weight
      # group.
      self.assertEqual(1, len(groups))
      self.assertEqual(3, len(groups[0]))
      kernel1 = groups[0][0]
      self.assertEqual('MyDense10/kernel', kernel1['name'])
      self.assertEqual('float32', kernel1['data'].dtype)
      self.assertEqual((3, 4), kernel1['data'].shape)
      self.assertTrue(np.allclose(np.ones([3, 4]), kernel1['data']))
      bias1 = groups[0][1]
      self.assertEqual('MyDense10/bias', bias1['name'])
      self.assertEqual('float32', bias1['data'].dtype)
      self.assertEqual((4,), bias1['data'].shape)
      self.assertTrue(np.allclose(np.zeros([4]), bias1['data']))
      kernel2 = groups[0][2]
      self.assertEqual('MyDense20/kernel', kernel2['name'])
      self.assertEqual('float32', kernel2['data'].dtype)
      self.assertEqual((4, 2), kernel2['data'].shape)
      self.assertTrue(np.allclose(np.ones([4, 2]), kernel2['data']))
  
    def testConvertWeightsFromSimpleModelSplitByLayer(self):
      input_tensor = tf.keras.layers.Input((3,))
      dense1 = tf.keras.layers.Dense(
          4, use_bias=True, kernel_initializer='ones', bias_initializer='zeros',
          name='MyDense30')(input_tensor)
      output = tf.keras.layers.Dense(
          2, use_bias=False, kernel_initializer='ones', name='MyDense40')(dense1)
      model = tf.keras.models.Model(inputs=[input_tensor], outputs=[output])
      h5_path = os.path.join(self._tmp_dir, 'MyModel.h5')
      model.save_weights(h5_path)
  
      # Load the saved weights as a JSON string.
      groups = conversion.h5_weights_to_tfjs_format(h5py.File(h5_path),
                                                    split_by_layer=True)
  
      # Check the loaded weights.
      # Due to `split_by_layer=True` and the fact that the model has two layers,
      # there should be two weight groups.
      self.assertEqual(2, len(groups))
      self.assertEqual(2, len(groups[0]))
      kernel1 = groups[0][0]
      self.assertEqual('MyDense30/kernel', kernel1['name'])
      self.assertEqual('float32', kernel1['data'].dtype)
      self.assertEqual((3, 4), kernel1['data'].shape)
      self.assertTrue(np.allclose(np.ones([3, 4]), kernel1['data']))
      bias1 = groups[0][1]
      self.assertEqual('MyDense30/bias', bias1['name'])
      self.assertEqual('float32', bias1['data'].dtype)
      self.assertEqual((4,), bias1['data'].shape)
      self.assertTrue(np.allclose(np.zeros([4]), bias1['data']))
  
      self.assertEqual(1, len(groups[1]))
      kernel2 = groups[1][0]
      self.assertEqual('MyDense40/kernel', kernel2['name'])
      self.assertEqual('float32', kernel2['data'].dtype)
      self.assertEqual((4, 2), kernel2['data'].shape)
      self.assertTrue(np.allclose(np.ones([4, 2]), kernel2['data']))
  
    def testConvertModelWithNestedLayerNames(self):
      model = tf.keras.Sequential()
  
      # Add a layer with a nested layer name, i.e., a layer name with slash(es)
      # in it.
      model.add(tf.keras.layers.Dense(2, input_shape=[12], name='dense'))
      model.add(tf.keras.layers.Dense(8, name='foo/dense'))
      model.add(tf.keras.layers.Dense(4, name='foo/bar/dense'))
      tfjs_path = os.path.join(self._tmp_dir, 'nested_layer_names_model')
      conversion.save_keras_model(model, tfjs_path)
  
      # Check model.json and weights manifest.
      with open(os.path.join(tfjs_path, 'model.json'), 'rt') as f:
        model_json = json.load(f)
  
      # Check meta-data in the artifact JSON.
      self.assertEqual(model_json['format'], 'layers-model')
      self.assertEqual(model_json['generatedBy'],
                       'keras v%s' % tf.keras.__version__)
      self.assertEqual(
          model_json['convertedBy'],
          'TensorFlow.js Converter v%s' % version.version)
  
      self.assertTrue(model_json['modelTopology'])
      weights_manifest = model_json['weightsManifest']
      weight_shapes = dict()
      for group in weights_manifest:
        for weight in group['weights']:
          weight_shapes[weight['name']] = weight['shape']
      self.assertEqual(
          sorted(['dense/kernel', 'dense/bias', 'foo/dense/kernel',
                  'foo/dense/bias', 'foo/bar/dense/kernel',
                  'foo/bar/dense/bias']),
          sorted(list(weight_shapes.keys())))
      self.assertEqual([12, 2], weight_shapes['dense/kernel'])
      self.assertEqual([2], weight_shapes['dense/bias'])
      self.assertEqual([2, 8], weight_shapes['foo/dense/kernel'])
      self.assertEqual([8], weight_shapes['foo/dense/bias'])
      self.assertEqual([8, 4], weight_shapes['foo/bar/dense/kernel'])
      self.assertEqual([4], weight_shapes['foo/bar/dense/bias'])
  
    def testConvertMergedModelFromSimpleModelNoSplitByLayer(self):
      input_tensor = tf.keras.layers.Input((3,))
      dense1 = tf.keras.layers.Dense(
          4, use_bias=True, kernel_initializer='ones', bias_initializer='zeros',
          name='MergedDense10')(input_tensor)
      output = tf.keras.layers.Dense(
          2, use_bias=False,
          kernel_initializer='ones', name='MergedDense20')(dense1)
      model = tf.keras.models.Model(inputs=[input_tensor], outputs=[output])
      h5_path = os.path.join(self._tmp_dir, 'MyModelMerged.h5')
      model.save(h5_path)
      config_json = json.loads(model.to_json(), encoding='utf8')
  
      # Load the saved weights as a JSON string.
      out, groups = conversion.h5_merged_saved_model_to_tfjs_format(
          h5py.File(h5_path))
      saved_topology = out['model_config']
  
      # check the model topology was stored
      self.assertEqual(config_json['class_name'], saved_topology['class_name'])
      self.assertEqual(config_json['config'], saved_topology['config'])
  
      # Check the loaded weights.
      # By default, all weights of the model ought to be put in the same group.
      self.assertEqual(1, len(groups))
  
      self.assertEqual(tf.keras.__version__, out['keras_version'])
      self.assertEqual('tensorflow', out['backend'])
      weight_group = groups[0]
      self.assertEqual(3, len(weight_group))
      kernel1 = weight_group[0]
      self.assertEqual('MergedDense10/kernel', kernel1['name'])
      self.assertEqual('float32', kernel1['data'].dtype)
      self.assertEqual((3, 4), kernel1['data'].shape)
      self.assertTrue(np.allclose(np.ones([3, 4]), kernel1['data']))
      bias1 = weight_group[1]
      self.assertEqual('MergedDense10/bias', bias1['name'])
      self.assertEqual('float32', bias1['data'].dtype)
      self.assertEqual((4,), bias1['data'].shape)
      self.assertTrue(np.allclose(np.zeros([4]), bias1['data']))
      kernel2 = weight_group[2]
      self.assertEqual('MergedDense20/kernel', kernel2['name'])
      self.assertEqual('float32', kernel2['data'].dtype)
      self.assertEqual((4, 2), kernel2['data'].shape)
      self.assertTrue(np.allclose(np.ones([4, 2]), kernel2['data']))
  
    def testConvertMergedModelFromSimpleModelSplitByLayer(self):
      input_tensor = tf.keras.layers.Input((3,))
      dense1 = tf.keras.layers.Dense(
          4, use_bias=True, kernel_initializer='ones', bias_initializer='zeros',
          name='MergedDense30')(input_tensor)
      output = tf.keras.layers.Dense(
          2, use_bias=False,
          kernel_initializer='ones', name='MergedDense40')(dense1)
      model = tf.keras.models.Model(inputs=[input_tensor], outputs=[output])
      h5_path = os.path.join(self._tmp_dir, 'MyModelMerged.h5')
      model.save(h5_path)
      config_json = json.loads(model.to_json(), encoding='utf8')
  
      # Load the saved weights as a JSON string.
      out, groups = conversion.h5_merged_saved_model_to_tfjs_format(
          h5py.File(h5_path), split_by_layer=True)
      saved_topology = out['model_config']
  
      # check the model topology was stored
      self.assertEqual(config_json['class_name'], saved_topology['class_name'])
      self.assertEqual(config_json['config'], saved_topology['config'])
  
      # Check the loaded weights.
      # Due to `split_by_layer=True`, there ought to be two weight groups,
      # because the model has two layers.
      self.assertEqual(2, len(groups))
  
      self.assertEqual(tf.keras.__version__, out['keras_version'])
      self.assertEqual('tensorflow', out['backend'])
      self.assertEqual(2, len(groups[0]))
      kernel1 = groups[0][0]
      self.assertEqual('MergedDense30/kernel', kernel1['name'])
      self.assertEqual('float32', kernel1['data'].dtype)
      self.assertEqual((3, 4), kernel1['data'].shape)
      self.assertTrue(np.allclose(np.ones([3, 4]), kernel1['data']))
      bias1 = groups[0][1]
      self.assertEqual('MergedDense30/bias', bias1['name'])
      self.assertEqual('float32', bias1['data'].dtype)
      self.assertEqual((4,), bias1['data'].shape)
      self.assertTrue(np.allclose(np.zeros([4]), bias1['data']))
      self.assertEqual(1, len(groups[1]))
      kernel2 = groups[1][0]
      self.assertEqual('MergedDense40/kernel', kernel2['name'])
      self.assertEqual('float32', kernel2['data'].dtype)
      self.assertEqual((4, 2), kernel2['data'].shape)
      self.assertTrue(np.allclose(np.ones([4, 2]), kernel2['data']))
  
    def testConvertWeightsFromSequentialModelNoSplitByLayer(self):
      sequential_model = tf.keras.models.Sequential([
          tf.keras.layers.Dense(
              3, input_shape=(2,), use_bias=True, kernel_initializer='ones',
              name='Dense10'),
          tf.keras.layers.Dense(
              1, use_bias=False, kernel_initializer='ones', name='Dense20')])
      h5_path = os.path.join(self._tmp_dir, 'SequentialModel.h5')
      sequential_model.save_weights(h5_path)
  
      # Load the saved weights as a JSON string.
      groups = conversion.h5_weights_to_tfjs_format(h5py.File(h5_path))
  
      # Check the loaded weights.
      # Due to the default `split_by_layer=False`, there should be only one weight
      # group.
      self.assertEqual(1, len(groups))
      self.assertEqual(3, len(groups[0]))
      kernel1 = groups[0][0]
      self.assertEqual('Dense10/kernel', kernel1['name'])
      self.assertEqual('float32', kernel1['data'].dtype)
      self.assertEqual((2, 3), kernel1['data'].shape)
      self.assertTrue(np.allclose(np.ones([2, 3]).tolist(), kernel1['data']))
      bias1 = groups[0][1]
      self.assertEqual('Dense10/bias', bias1['name'])
      self.assertEqual('float32', bias1['data'].dtype)
      self.assertEqual((3,), bias1['data'].shape)
      self.assertTrue(np.allclose(np.zeros([3]).tolist(), bias1['data']))
      kernel2 = groups[0][2]
      self.assertEqual('Dense20/kernel', kernel2['name'])
      self.assertEqual('float32', kernel2['data'].dtype)
      self.assertEqual((3, 1), kernel2['data'].shape)
      self.assertTrue(np.allclose(np.ones([3, 1]).tolist(), kernel2['data']))
  
    def testConvertWeightsFromSequentialModelSplitByLayer(self):
      sequential_model = tf.keras.models.Sequential([
          tf.keras.layers.Dense(
              3, input_shape=(2,), use_bias=True, kernel_initializer='ones',
              name='Dense30'),
          tf.keras.layers.Dense(
              1, use_bias=False, kernel_initializer='ones', name='Dense40')])
      h5_path = os.path.join(self._tmp_dir, 'SequentialModel.h5')
      sequential_model.save_weights(h5_path)
  
      # Load the saved weights as a JSON string.
      groups = conversion.h5_weights_to_tfjs_format(h5py.File(h5_path),
                                                    split_by_layer=True)
  
      # Check the loaded weights.
      # Due to the default `split_by_layer=True`, there should be two weight
      # gropus, because the model has two layers.
      self.assertEqual(2, len(groups))
      self.assertEqual(2, len(groups[0]))
      kernel1 = groups[0][0]
      self.assertEqual('Dense30/kernel', kernel1['name'])
      self.assertEqual('float32', kernel1['data'].dtype)
      self.assertEqual((2, 3), kernel1['data'].shape)
      self.assertTrue(np.allclose(np.ones([2, 3]).tolist(), kernel1['data']))
      bias1 = groups[0][1]
      self.assertEqual('Dense30/bias', bias1['name'])
      self.assertEqual('float32', bias1['data'].dtype)
      self.assertEqual((3,), bias1['data'].shape)
      self.assertTrue(np.allclose(np.zeros([3]).tolist(), bias1['data']))
  
      self.assertEqual(1, len(groups[1]))
      kernel2 = groups[1][0]
      self.assertEqual('Dense40/kernel', kernel2['name'])
      self.assertEqual('float32', kernel2['data'].dtype)
      self.assertEqual((3, 1), kernel2['data'].shape)
      self.assertTrue(np.allclose(np.ones([3, 1]).tolist(), kernel2['data']))
  
    def testSaveModelSucceedsForNonSequentialModel(self):
      t_input = tf.keras.Input([2])
      dense_layer = tf.keras.layers.Dense(3)
      t_output = dense_layer(t_input)
      model = tf.keras.Model(t_input, t_output)
      conversion.save_keras_model(model, self._tmp_dir)
  
      # Verify the content of the artifacts output directory.
      self.assertTrue(
          os.path.isfile(os.path.join(self._tmp_dir, 'group1-shard1of1.bin')))
      model_json = json.load(
          open(os.path.join(self._tmp_dir, 'model.json'), 'rt'))
  
      topology_json = model_json['modelTopology']
      self.assertIn('keras_version', topology_json)
      self.assertIn('backend', topology_json)
      self.assertIn('model_config', topology_json)
  
      weights_manifest = model_json['weightsManifest']
      self.assertIsInstance(weights_manifest, list)
      self.assertEqual(1, len(weights_manifest))
      self.assertIn('paths', weights_manifest[0])
  
    def testSaveModelSucceedsForTfKerasNonSequentialModel(self):
      t_input = tf.keras.Input([2])
      dense_layer = tf.keras.layers.Dense(3)
      t_output = dense_layer(t_input)
      model = tf.keras.Model(t_input, t_output)
  
      # `tf.keras.Model`s must be compiled before they can be saved.
      model.compile(loss='mean_squared_error', optimizer='sgd')
  
      conversion.save_keras_model(model, self._tmp_dir)
  
      # Verify the content of the artifacts output directory.
      self.assertTrue(
          os.path.isfile(os.path.join(self._tmp_dir, 'group1-shard1of1.bin')))
      model_json = json.load(
          open(os.path.join(self._tmp_dir, 'model.json'), 'rt'))
  
      topology_json = model_json['modelTopology']
      self.assertIn('keras_version', topology_json)
      self.assertIn('backend', topology_json)
      self.assertIn('model_config', topology_json)
  
      weights_manifest = model_json['weightsManifest']
      self.assertIsInstance(weights_manifest, list)
      self.assertEqual(1, len(weights_manifest))
      self.assertIn('paths', weights_manifest[0])
  
    def testSaveModelSucceedsForNestedKerasModel(self):
      inner_model = tf.keras.Sequential([
          tf.keras.layers.Dense(4, input_shape=[3], activation='relu'),
          tf.keras.layers.Dense(3, activation='tanh')])
      outer_model = tf.keras.Sequential()
      outer_model.add(inner_model)
      outer_model.add(tf.keras.layers.Dense(1, activation='sigmoid'))
  
      conversion.save_keras_model(outer_model, self._tmp_dir)
  
      # Verify the content of the artifacts output directory.
      self.assertTrue(
          os.path.isfile(os.path.join(self._tmp_dir, 'group1-shard1of1.bin')))
      model_json = json.load(
          open(os.path.join(self._tmp_dir, 'model.json'), 'rt'))
  
      topology_json = model_json['modelTopology']
      self.assertIn('keras_version', topology_json)
      self.assertIn('backend', topology_json)
      self.assertIn('model_config', topology_json)
  
      # Verify that all the layers' weights are present.
      weights_manifest = model_json['weightsManifest']
      self.assertIsInstance(weights_manifest, list)
      weight_entries = []
      for group in weights_manifest:
        weight_entries.extend(group['weights'])
      self.assertEqual(6, len(weight_entries))
  
    def testSaveModelSucceedsForTfKerasSequentialModel(self):
      model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=[2])])
  
      # `tf.keras.Model`s must be compiled before they can be saved.
      model.compile(loss='mean_squared_error', optimizer='sgd')
  
      conversion.save_keras_model(model, self._tmp_dir)
  
      # Verify the content of the artifacts output directory.
      self.assertTrue(
          os.path.isfile(os.path.join(self._tmp_dir, 'group1-shard1of1.bin')))
      model_json = json.load(
          open(os.path.join(self._tmp_dir, 'model.json'), 'rt'))
  
      topology_json = model_json['modelTopology']
      self.assertIn('keras_version', topology_json)
      self.assertIn('backend', topology_json)
      self.assertIn('model_config', topology_json)
  
      weights_manifest = model_json['weightsManifest']
      self.assertIsInstance(weights_manifest, list)
      self.assertEqual(1, len(weights_manifest))
      self.assertIn('paths', weights_manifest[0])
  
    def testSavedModelSucceedsForExistingDirAndSequential(self):
      artifacts_dir = os.path.join(self._tmp_dir, 'artifacts')
      os.makedirs(artifacts_dir)
      model = tf.keras.Sequential()
      model.add(tf.keras.layers.Dense(3, input_shape=[2]))
      conversion.save_keras_model(model, artifacts_dir)
  
      # Verify the content of the artifacts output directory.
      self.assertTrue(
          os.path.isfile(os.path.join(artifacts_dir, 'group1-shard1of1.bin')))
      model_json = json.load(
          open(os.path.join(artifacts_dir, 'model.json'), 'rt'))
  
      topology_json = model_json['modelTopology']
      self.assertIn('keras_version', topology_json)
      self.assertIn('backend', topology_json)
      self.assertIn('model_config', topology_json)
  
      weights_manifest = model_json['weightsManifest']
      self.assertIsInstance(weights_manifest, list)
      self.assertEqual(1, len(weights_manifest))
      self.assertIn('paths', weights_manifest[0])
  
    def testSavedModelSucceedsForCustomShardSize(self):
      model = tf.keras.Sequential([
          tf.keras.layers.Dense(1, input_shape=[2], activation='relu'),
          tf.keras.layers.Dense(3, activation='tanh')
      ])
  
      weights = model.get_weights()
      total_weight_bytes = sum(np.size(w) for w in weights) * 4
  
      # Due to the shard size, there ought to be 4 shards after conversion.
      weight_shard_size_bytes = int(total_weight_bytes * 0.3)
  
      # Convert Keras model to tfjs_layers_model format.
      conversion.save_keras_model(model, self._tmp_dir,
                                  weight_shard_size_bytes=weight_shard_size_bytes)
  
      weight_files = sorted(glob.glob(os.path.join(self._tmp_dir, 'group*.bin')))
      self.assertEqual(len(weight_files), 4)
      weight_file_sizes = [os.path.getsize(f) for f in weight_files]
      self.assertEqual(sum(weight_file_sizes), total_weight_bytes)
      self.assertEqual(weight_file_sizes[0], weight_file_sizes[1])
      self.assertEqual(weight_file_sizes[0], weight_file_sizes[2])
      self.assertLess(weight_file_sizes[3], weight_file_sizes[0])
  
    def testSavedModelRaisesErrorIfArtifactsDirExistsAsAFile(self):
      artifacts_dir = os.path.join(self._tmp_dir, 'artifacts')
      with open(artifacts_dir, 'wt') as f:
        f.write('foo\n')
      t_input = tf.keras.Input([2])
      dense_layer = tf.keras.layers.Dense(3)
      t_output = dense_layer(t_input)
      model = tf.keras.Model(t_input, t_output)
      with self.assertRaisesRegexp(  # pylint: disable=deprecated-method
          ValueError, r'already exists as a file'):
        conversion.save_keras_model(model, artifacts_dir)
  
    def testTranslateBatchNormalizationV1ClassName(self):
      # The config JSON of a model consisting of a "BatchNormalizationV1" layer.
      # pylint: disable=line-too-long
      json_object = json.loads(
          '{"class_name": "Sequential", "keras_version": "2.2.4-tf", "config": {"layers": [{"class_name": "Dense", "config": {"kernel_initializer": {"class_name": "GlorotUniform", "config": {"dtype": "float32", "seed": null}}, "name": "dense", "kernel_constraint": null, "bias_regularizer": null, "bias_constraint": null, "dtype": "float32", "activation": "relu", "trainable": true, "kernel_regularizer": null, "bias_initializer": {"class_name": "Zeros", "config": {"dtype": "float32"}}, "units": 10, "batch_input_shape": [null, 3], "use_bias": true, "activity_regularizer": null}}, {"class_name": "BatchNormalizationV1", "config": {"beta_constraint": null, "gamma_initializer": {"class_name": "Ones", "config": {"dtype": "float32"}}, "moving_mean_initializer": {"class_name": "Zeros", "config": {"dtype": "float32"}}, "name": "batch_normalization_v1", "dtype": "float32", "trainable": true, "moving_variance_initializer": {"class_name": "Ones", "config": {"dtype": "float32"}}, "beta_initializer": {"class_name": "Zeros", "config": {"dtype": "float32"}}, "scale": true, "axis": [1], "epsilon": 0.001, "gamma_constraint": null, "gamma_regularizer": null, "beta_regularizer": null, "momentum": 0.99, "center": true}}, {"class_name": "Dense", "config": {"kernel_initializer": {"class_name": "GlorotUniform", "config": {"dtype": "float32", "seed": null}}, "name": "dense_1", "kernel_constraint": null, "bias_regularizer": null, "bias_constraint": null, "dtype": "float32", "activation": "linear", "trainable": true, "kernel_regularizer": null, "bias_initializer": {"class_name": "Zeros", "config": {"dtype": "float32"}}, "units": 1, "use_bias": true, "activity_regularizer": null}}], "name": "sequential"}, "backend": "tensorflow"}')
      # pylint: enable=line-too-long
      conversion.translate_class_names(json_object)
      # Some class names should not have been changed be translate_class_names().
      self.assertEqual(json_object['class_name'], 'Sequential')
      self.assertEqual(json_object['keras_version'], '2.2.4-tf')
      self.assertEqual(json_object['config']['layers'][0]['class_name'], 'Dense')
      # The translation should have happend:
      #     BatchNormalizationV1 --> BatchNormalization.
      self.assertEqual(
          json_object['config']['layers'][1]['class_name'], 'BatchNormalization')
      self.assertEqual(json_object['config']['layers'][2]['class_name'], 'Dense')
  
      # Assert that converted JSON can be reconstituted as a model object.
      model = tf.keras.models.model_from_json(json.dumps(json_object))
      self.assertIsInstance(model, tf.keras.Sequential)
      self.assertEqual(model.input_shape, (None, 3))
      self.assertEqual(model.output_shape, (None, 1))
      self.assertEqual(model.layers[0].units, 10)
      self.assertEqual(model.layers[2].units, 1)
  
    def testTranslateUnifiedGRUAndLSTMClassName(self):
      # The config JSON of a model consisting of a "UnifiedGRU" layer
      # and a "UnifiedLSTM" layer.
      # pylint: disable=line-too-long
      json_object = json.loads(
          '{"class_name": "Sequential", "keras_version": "2.2.4-tf", "config": {"layers": [{"class_name": "UnifiedGRU", "config": {"recurrent_activation": "sigmoid", "dtype": "float32", "trainable": true, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"seed": null, "gain": 1.0}}, "use_bias": true, "bias_regularizer": null, "return_state": false, "unroll": false, "activation": "tanh", "bias_initializer": {"class_name": "Zeros", "config": {}}, "units": 10, "batch_input_shape": [null, 4, 3], "activity_regularizer": null, "recurrent_dropout": 0.0, "kernel_initializer": {"class_name": "GlorotUniform", "config": {"seed": null}}, "kernel_constraint": null, "time_major": false, "dropout": 0.0, "stateful": false, "reset_after": true, "recurrent_regularizer": null, "name": "unified_gru", "bias_constraint": null, "go_backwards": false, "implementation": 1, "kernel_regularizer": null, "return_sequences": true, "recurrent_constraint": null}}, {"class_name": "UnifiedLSTM", "config": {"recurrent_activation": "sigmoid", "dtype": "float32", "trainable": true, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"seed": null, "gain": 1.0}}, "use_bias": true, "bias_regularizer": null, "return_state": false, "unroll": false, "activation": "tanh", "bias_initializer": {"class_name": "Zeros", "config": {}}, "units": 2, "unit_forget_bias": true, "activity_regularizer": null, "recurrent_dropout": 0.0, "kernel_initializer": {"class_name": "GlorotUniform", "config": {"seed": null}}, "kernel_constraint": null, "time_major": false, "dropout": 0.0, "stateful": false, "recurrent_regularizer": null, "name": "unified_lstm", "bias_constraint": null, "go_backwards": false, "implementation": 1, "kernel_regularizer": null, "return_sequences": false, "recurrent_constraint": null}}], "name": "sequential"}, "backend": "tensorflow"}')
      # pylint: enable=line-too-long
      conversion.translate_class_names(json_object)
      # Some class names should not have been changed be translate_class_names().
      self.assertEqual(json_object['class_name'], 'Sequential')
      self.assertEqual(json_object['keras_version'], '2.2.4-tf')
      # The translation should have happend:
      #     UnifiedGRU --> GRU.
      #     UnifiedLSTM --> LSTM.
      self.assertEqual(json_object['config']['layers'][0]['class_name'], 'GRU')
      self.assertEqual(json_object['config']['layers'][1]['class_name'], 'LSTM')
  
      # Assert that converted JSON can be reconstituted as a model object.
      model = tf.keras.models.model_from_json(json.dumps(json_object))
      self.assertIsInstance(model, tf.keras.Sequential)
      self.assertEqual(model.input_shape, (None, 4, 3))
      self.assertEqual(model.output_shape, (None, 2))
  
  
  if __name__ == '__main__':
    tf.test.main()
  