text: # Copyright 2020 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 Keras discretization preprocessing layer."""
  
  import tensorflow.compat.v2 as tf
  
  import os
  
  from absl.testing import parameterized
  
  import numpy as np
  
  import keras
  from keras import keras_parameterized
  from keras import testing_utils
  from keras.layers.preprocessing import discretization
  from keras.layers.preprocessing import preprocessing_test_utils
  
  
  @keras_parameterized.run_all_keras_modes
  class DiscretizationTest(keras_parameterized.TestCase,
                           preprocessing_test_utils.PreprocessingLayerTest):
  
    def test_bucketize_with_explicit_buckets_integer(self):
      input_array = np.array([[-1.5, 1.0, 3.4, .5], [0.0, 3.0, 1.3, 0.0]])
  
      expected_output = [[0, 1, 3, 1], [0, 3, 2, 0]]
      expected_output_shape = [None, 4]
  
      input_data = keras.Input(shape=(4,))
      layer = discretization.Discretization(bin_boundaries=[0., 1., 2.])
      bucket_data = layer(input_data)
      self.assertAllEqual(expected_output_shape, bucket_data.shape.as_list())
  
      model = keras.Model(inputs=input_data, outputs=bucket_data)
      output_dataset = model.predict(input_array)
      self.assertAllEqual(expected_output, output_dataset)
  
    def test_bucketize_with_explicit_buckets_int_input(self):
      input_array = np.array([[-1, 1, 3, 0], [0, 3, 1, 0]], dtype=np.int64)
  
      expected_output = [[0, 2, 3, 1], [1, 3, 2, 1]]
      expected_output_shape = [None, 4]
  
      input_data = keras.Input(shape=(4,), dtype=tf.int64)
      layer = discretization.Discretization(bin_boundaries=[-.5, 0.5, 1.5])
      bucket_data = layer(input_data)
      self.assertAllEqual(expected_output_shape, bucket_data.shape.as_list())
  
      model = keras.Model(inputs=input_data, outputs=bucket_data)
      output_dataset = model.predict(input_array)
      self.assertAllEqual(expected_output, output_dataset)
  
    def test_bucketize_with_explicit_buckets_sparse_float_input(self):
      indices = [[0, 1], [0, 2], [1, 1]]
      input_array = tf.SparseTensor(
          indices=indices, values=[-1.5, 1.0, 3.4], dense_shape=[2, 3])
      expected_output = [0, 2, 3]
      input_data = keras.Input(shape=(3,), dtype=tf.float32, sparse=True)
      layer = discretization.Discretization(bin_boundaries=[-.5, 0.5, 1.5])
      bucket_data = layer(input_data)
  
      model = keras.Model(inputs=input_data, outputs=bucket_data)
      output_dataset = model.predict(input_array, steps=1)
      self.assertAllEqual(indices, output_dataset.indices)
      self.assertAllEqual(expected_output, output_dataset.values)
  
    def test_bucketize_with_explicit_buckets_ragged_float_input(self):
      input_array = tf.ragged.constant([[-1.5, 1.0, 3.4, .5],
                                                 [0.0, 3.0, 1.3]])
  
      expected_output = [[0, 1, 3, 1], [0, 3, 2]]
      expected_output_shape = [None, None]
  
      input_data = keras.Input(shape=(None,), ragged=True)
      layer = discretization.Discretization(bin_boundaries=[0., 1., 2.])
      bucket_data = layer(input_data)
      self.assertAllEqual(expected_output_shape, bucket_data.shape.as_list())
  
      model = keras.Model(inputs=input_data, outputs=bucket_data)
      output_dataset = model.predict(input_array)
      self.assertAllEqual(expected_output, output_dataset)
  
    def test_bucketize_with_explicit_buckets_ragged_int_input(self):
      input_array = tf.ragged.constant([[-1, 1, 3, 0], [0, 3, 1]],
                                                dtype=tf.int64)
  
      expected_output = [[0, 2, 3, 1], [1, 3, 2]]
      expected_output_shape = [None, None]
  
      input_data = keras.Input(shape=(None,), ragged=True, dtype=tf.int64)
      layer = discretization.Discretization(bin_boundaries=[-.5, 0.5, 1.5])
      bucket_data = layer(input_data)
      self.assertAllEqual(expected_output_shape, bucket_data.shape.as_list())
      model = keras.Model(inputs=input_data, outputs=bucket_data)
      output_dataset = model.predict(input_array)
      self.assertAllEqual(expected_output, output_dataset)
  
    def test_bucketize_with_explicit_buckets_sparse_int_input(self):
      indices = [[0, 1], [0, 2], [1, 1]]
      input_array = tf.SparseTensor(
          indices=indices, values=[-1, 1, 3], dense_shape=[2, 3])
      expected_output = [0, 2, 3]
      input_data = keras.Input(shape=(3,), dtype=tf.int32, sparse=True)
      layer = discretization.Discretization(bin_boundaries=[-.5, 0.5, 1.5])
      bucket_data = layer(input_data)
  
      model = keras.Model(inputs=input_data, outputs=bucket_data)
      output_dataset = model.predict(input_array, steps=1)
      self.assertAllEqual(indices, output_dataset.indices)
      self.assertAllEqual(expected_output, output_dataset.values)
  
    def test_output_shape(self):
      input_data = keras.Input(batch_size=16, shape=(4,), dtype=tf.string)
      layer = discretization.Discretization(bin_boundaries=[-.5, 0.5, 1.5])
      output = layer(input_data)
      self.assertAllEqual(output.shape.as_list(), [16, 4])
  
    def test_num_bins_negative_fails(self):
      with self.assertRaisesRegex(ValueError, "`num_bins` must be.*num_bins=-7"):
        _ = discretization.Discretization(num_bins=-7)
  
    def test_num_bins_and_bins_set_fails(self):
      with self.assertRaisesRegex(
          ValueError,
          r"`num_bins` and `bin_boundaries` should not be set.*5.*\[1, 2\]"):
        _ = discretization.Discretization(num_bins=5, bins=[1, 2])
  
  
  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
  class DiscretizationAdaptTest(keras_parameterized.TestCase,
                                preprocessing_test_utils.PreprocessingLayerTest):
  
    @parameterized.named_parameters([
        {
            "testcase_name": "2d_single_element",
            "adapt_data": np.array([[1.], [2.], [3.], [4.], [5.]]),
            "test_data": np.array([[1.], [2.], [3.]]),
            "use_dataset": True,
            "expected": np.array([[0], [1], [2]]),
            "num_bins": 5,
            "epsilon": 0.01
        }, {
            "testcase_name": "2d_multi_element",
            "adapt_data": np.array([[1., 6.], [2., 7.], [3., 8.], [4., 9.],
                                    [5., 10.]]),
            "test_data": np.array([[1., 10.], [2., 6.], [3., 8.]]),
            "use_dataset": True,
            "expected": np.array([[0, 4], [0, 2], [1, 3]]),
            "num_bins": 5,
            "epsilon": 0.01
        }, {
            "testcase_name": "1d_single_element",
            "adapt_data": np.array([3., 2., 1., 5., 4.]),
            "test_data": np.array([1., 2., 3.]),
            "use_dataset": True,
            "expected": np.array([0, 1, 2]),
            "num_bins": 5,
            "epsilon": 0.01
        }, {
            "testcase_name": "300_batch_1d_single_element_1",
            "adapt_data": np.arange(300),
            "test_data": np.arange(300),
            "use_dataset": True,
            "expected":
                np.concatenate([np.zeros(101), np.ones(99), 2 * np.ones(100)]),
            "num_bins": 3,
            "epsilon": 0.01
        }, {
            "testcase_name": "300_batch_1d_single_element_2",
            "adapt_data": np.arange(300) ** 2,
            "test_data": np.arange(300) ** 2,
            "use_dataset": True,
            "expected":
                np.concatenate([np.zeros(101), np.ones(99), 2 * np.ones(100)]),
            "num_bins": 3,
            "epsilon": 0.01
        }, {
            "testcase_name": "300_batch_1d_single_element_large_epsilon",
            "adapt_data": np.arange(300),
            "test_data": np.arange(300),
            "use_dataset": True,
            "expected": np.concatenate([np.zeros(137), np.ones(163)]),
            "num_bins": 2,
            "epsilon": 0.1
        }])
    def test_layer_computation(self, adapt_data, test_data, use_dataset,
                               expected, num_bins=5, epsilon=0.01):
  
      input_shape = tuple(list(test_data.shape)[1:])
      np.random.shuffle(adapt_data)
      if use_dataset:
        # Keras APIs expect batched datasets
        adapt_data = tf.data.Dataset.from_tensor_slices(adapt_data).batch(
            test_data.shape[0] // 2)
        test_data = tf.data.Dataset.from_tensor_slices(test_data).batch(
            test_data.shape[0] // 2)
  
      layer = discretization.Discretization(epsilon=epsilon, num_bins=num_bins)
      layer.adapt(adapt_data)
  
      input_data = keras.Input(shape=input_shape)
      output = layer(input_data)
      model = keras.Model(input_data, output)
      model._run_eagerly = testing_utils.should_run_eagerly()
      output_data = model.predict(test_data)
      self.assertAllClose(expected, output_data)
  
    def test_merge_state(self):
      data = np.arange(300)
      partial_ds_1 = tf.data.Dataset.from_tensor_slices(data[:100])
      partial_ds_2 = tf.data.Dataset.from_tensor_slices(data[100:200])
      partial_ds_3 = tf.data.Dataset.from_tensor_slices(data[200:])
      full_ds = partial_ds_1.concatenate(partial_ds_2).concatenate(partial_ds_3)
  
      # Use a higher epsilon to avoid any discrepencies from the quantile
      # approximation.
      full_layer = discretization.Discretization(num_bins=3, epsilon=0.001)
      full_layer.adapt(full_ds.batch(2))
  
      partial_layer_1 = discretization.Discretization(num_bins=3, epsilon=0.001)
      partial_layer_1.adapt(partial_ds_1.batch(2))
      partial_layer_2 = discretization.Discretization(num_bins=3, epsilon=0.001)
      partial_layer_2.adapt(partial_ds_2.batch(2))
      partial_layer_3 = discretization.Discretization(num_bins=3, epsilon=0.001)
      partial_layer_3.adapt(partial_ds_3.batch(2))
      partial_layer_1.merge_state([partial_layer_2, partial_layer_3])
      merged_layer = partial_layer_1
  
      data = np.arange(300)
      self.assertAllClose(full_layer(data), merged_layer(data))
  
    def test_merge_with_stateless_layers_fails(self):
      layer1 = discretization.Discretization(num_bins=2, name="layer1")
      layer1.adapt([1, 2, 3])
      layer2 = discretization.Discretization(bin_boundaries=[0, 1], name="layer2")
      with self.assertRaisesRegex(ValueError, "Cannot merge.*layer2"):
        layer1.merge_state([layer2])
  
    def test_merge_with_unadapted_layers_fails(self):
      layer1 = discretization.Discretization(num_bins=2, name="layer1")
      layer1.adapt([1, 2, 3])
      layer2 = discretization.Discretization(num_bins=2, name="layer2")
      with self.assertRaisesRegex(ValueError, "Cannot merge.*layer2"):
        layer1.merge_state([layer2])
  
    def test_multiple_adapts(self):
      first_adapt = [[1], [2], [3]]
      second_adapt = [[4], [5], [6]]
      predict_input = [[2], [2]]
      expected_first_output = [[1], [1]]
      expected_second_output = [[0], [0]]
  
      inputs = keras.Input(shape=(1,), dtype=tf.int32)
      layer = discretization.Discretization(num_bins=3)
      layer.adapt(first_adapt)
      outputs = layer(inputs)
      model = keras.Model(inputs=inputs, outputs=outputs)
  
      actual_output = model.predict(predict_input)
      self.assertAllClose(actual_output, expected_first_output)
  
      # Re-adapt the layer on new inputs.
      layer.adapt(second_adapt)
      # Re-compile the model.
      model.compile()
      # `predict` should now use the new model state.
      actual_output = model.predict(predict_input)
      self.assertAllClose(actual_output, expected_second_output)
  
    def test_saved_model_tf(self):
      input_data = [[1], [2], [3]]
      expected_output = [[0], [1], [2]]
  
      inputs = keras.Input(shape=(1,), dtype=tf.int32)
      layer = discretization.Discretization(num_bins=3)
      layer.adapt(input_data)
      outputs = layer(inputs)
      model = keras.Model(inputs=inputs, outputs=outputs)
  
      output_data = model.predict(input_data)
      self.assertAllClose(output_data, expected_output)
  
      # Save the model to disk.
      output_path = os.path.join(self.get_temp_dir(), "tf_saved_model")
      tf.saved_model.save(model, output_path)
      loaded_model = tf.saved_model.load(output_path)
      f = loaded_model.signatures["serving_default"]
  
      # Ensure that the loaded model is unique (so that the save/load is real)
      self.assertIsNot(model, loaded_model)
  
      # Validate correctness of the new model.
      new_output_data = f(tf.constant(input_data))["discretization"]
      self.assertAllClose(new_output_data, expected_output)
  
    def test_saved_model_keras(self):
      input_data = [[1], [2], [3]]
      expected_output = [[0], [1], [2]]
  
      cls = discretization.Discretization
      inputs = keras.Input(shape=(1,), dtype=tf.int32)
      layer = cls(num_bins=3)
      layer.adapt(input_data)
      outputs = layer(inputs)
      model = keras.Model(inputs=inputs, outputs=outputs)
  
      output_data = model.predict(input_data)
      self.assertAllClose(output_data, expected_output)
  
      # Save the model to disk.
      output_path = os.path.join(self.get_temp_dir(), "tf_keras_saved_model")
      model.save(output_path, save_format="tf")
      loaded_model = keras.models.load_model(
          output_path, custom_objects={"Discretization": cls})
  
      # Ensure that the loaded model is unique (so that the save/load is real)
      self.assertIsNot(model, loaded_model)
  
      # Validate correctness of the new model.
      new_output_data = loaded_model.predict(input_data)
      self.assertAllClose(new_output_data, expected_output)
  
    def test_saved_weights_keras(self):
      input_data = [[1], [2], [3]]
      expected_output = [[0], [1], [2]]
  
      cls = discretization.Discretization
      inputs = keras.Input(shape=(1,), dtype=tf.int32)
      layer = cls(num_bins=3)
      layer.adapt(input_data)
      outputs = layer(inputs)
      model = keras.Model(inputs=inputs, outputs=outputs)
  
      output_data = model.predict(input_data)
      self.assertAllClose(output_data, expected_output)
  
      # Save the model to disk.
      output_path = os.path.join(self.get_temp_dir(), "tf_keras_saved_weights")
      model.save_weights(output_path, save_format="tf")
      new_model = keras.Model.from_config(
          model.get_config(), custom_objects={"Discretization": cls})
      new_model.load_weights(output_path)
  
      # Validate correctness of the new model.
      new_output_data = new_model.predict(input_data)
      self.assertAllClose(new_output_data, expected_output)
  
  
  if __name__ == "__main__":
    tf.test.main()
  