text: # Copyright 2017 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.
  # ==============================================================================
  """Functional tests for the ops to generate and execute vocab remapping."""
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import os
  import numpy as np
  
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import gen_checkpoint_ops
  from tensorflow.python.ops import partitioned_variables
  from tensorflow.python.ops import variable_scope
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import flags
  from tensorflow.python.platform import test
  from tensorflow.python.training import saver
  
  FLAGS = flags.FLAGS
  
  
  class GenerateVocabRemappingTest(test.TestCase):
    """Tests for the generate_vocab_remapping() method."""
  
    def setUp(self):
      self.new_vocab_file = os.path.join(self.get_temp_dir(),
                                         'keyword_shifted.txt')
      with open(self.new_vocab_file, 'w') as f:
        f.write('\n'.join(['MISSING', 'knitting', 'eminem']) + '\n')
      self.old_vocab_file = os.path.join(self.get_temp_dir(),
                                         'keyword.txt')
      with open(self.old_vocab_file, 'w') as f:
        f.write('\n'.join(['knitting', 'eminem', 'MISSING']) + '\n')
  
    def test_generate_remapping_with_no_vocab_changes(self):
      """Tests where vocab does not change at all."""
      remapping, num_present = gen_checkpoint_ops._generate_vocab_remapping(
          new_vocab_file=self.old_vocab_file,
          old_vocab_file=self.old_vocab_file,
          num_new_vocab=3,
          new_vocab_offset=0)
      expected_remapping = range(0, 3)
      expected_num_present = 3
      with self.test_session():
        self.assertAllEqual(expected_remapping, remapping.eval())
        self.assertAllEqual(expected_num_present, num_present.eval())
  
    def test_generate_remapping_with_shifted_vocab(self):
      """Tests where vocab is the same, but shifted / ordered differently."""
      remapping, num_present = gen_checkpoint_ops._generate_vocab_remapping(
          new_vocab_file=self.new_vocab_file,
          old_vocab_file=self.old_vocab_file,
          num_new_vocab=3,
          new_vocab_offset=0)
      expected_remapping = [2, 0, 1]
      expected_num_present = 3
      with self.test_session():
        self.assertAllEqual(expected_remapping, remapping.eval())
        self.assertAllEqual(expected_num_present, num_present.eval())
  
    def test_generate_remapping_with_offset(self):
      """Tests offset and num_new_vocab logic."""
      remapping, num_present = gen_checkpoint_ops._generate_vocab_remapping(
          new_vocab_file=self.new_vocab_file,
          old_vocab_file=self.old_vocab_file,
          num_new_vocab=1,
          new_vocab_offset=1)
      expected_remapping = [0]
      expected_num_present = 1
      with self.test_session():
        self.assertAllEqual(expected_remapping, remapping.eval())
        self.assertAllEqual(expected_num_present, num_present.eval())
  
  
  class LoadAndRemapMatrixTest(test.TestCase):
    """Tests for the load_and_remap_matrix() op."""
  
    def setUp(self):
      ops.reset_default_graph()
      self.old_num_rows = 5
      self.old_num_cols = 16
      self.matrix_value = np.reshape(
          range(0, self.old_num_rows * self.old_num_cols), (self.old_num_rows,
                                                            self.old_num_cols))
      with variable_scope.variable_scope('some_scope'):
        matrix = variable_scope.get_variable(
            'matrix',
            dtype=dtypes.float32,
            initializer=constant_op.constant(
                self.matrix_value, dtype=dtypes.float32))
        self.old_tensor_name = 'some_scope/matrix'
  
      save = saver.Saver([matrix])
      with self.test_session() as sess:
        variables.global_variables_initializer().run()
        self.bundle_file = os.path.join(test.get_temp_dir(), 'bundle_checkpoint')
        save.save(sess, self.bundle_file)
  
    def test_load_and_remap_no_missing(self):
      """Tests the op's load and remap where there are no missing entries."""
  
      # No column remapping, new weight matrix has second row, then first row.
      row_remapping = [1, 0]
      remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix(
          ckpt_path=[self.bundle_file],
          old_tensor_name=self.old_tensor_name,
          row_remapping=row_remapping,
          col_remapping=[],
          initializing_values=[],
          num_rows=2,
          num_cols=self.old_num_cols)
      with self.test_session():
        self.assertAllClose(self.matrix_value[row_remapping],
                            remapped_matrix.eval())
  
      # No row remapping, new weight matrix has third col, then first col.
      row_remapping = list(range(self.old_num_rows))
      col_remapping = [2, 0]
      remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix(
          ckpt_path=[self.bundle_file],
          old_tensor_name=self.old_tensor_name,
          row_remapping=row_remapping,
          col_remapping=col_remapping,
          initializing_values=[],
          num_rows=len(row_remapping),
          num_cols=len(col_remapping))
      with self.test_session():
        self.assertAllClose(self.matrix_value[row_remapping][:, col_remapping],
                            remapped_matrix.eval())
  
      # Both row and column remappings.
      row_remapping = [1, 0, 4]
      col_remapping = [1, 15]
      remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix(
          ckpt_path=[self.bundle_file],
          old_tensor_name=self.old_tensor_name,
          row_remapping=row_remapping,
          col_remapping=col_remapping,
          initializing_values=[],
          num_rows=len(row_remapping),
          num_cols=len(col_remapping))
      with self.test_session():
        self.assertAllClose(self.matrix_value[row_remapping][:, col_remapping],
                            remapped_matrix.eval())
  
    def test_load_and_remap_with_init(self):
      """Tests the op's load and remap where there are missing entries."""
      init_val = 42
      remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix(
          ckpt_path=[self.bundle_file],
          old_tensor_name=self.old_tensor_name,
          row_remapping=[2, -1, 0],
          col_remapping=[1, -1],
          initializing_values=[init_val] * 4,
          num_rows=3,
          num_cols=2)
  
      expected_remapped_matrix = np.reshape(
          [33, init_val, init_val, init_val, 1, init_val], [3, 2])
  
      with self.test_session():
        self.assertAllClose(expected_remapped_matrix, remapped_matrix.eval())
  
    def test_load_and_remap_all_missing_rows(self):
      """Tests when all the rows are missing and need to be initialized."""
      num_rows = 7
      initializing_values = [42] * num_rows * self.old_num_cols
      remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix(
          ckpt_path=[self.bundle_file],
          old_tensor_name=self.old_tensor_name,
          row_remapping=[-1] * num_rows,
          col_remapping=[],
          initializing_values=initializing_values,
          num_rows=num_rows,
          num_cols=self.old_num_cols)
      with self.test_session():
        self.assertAllClose(
            np.reshape(initializing_values, (num_rows, self.old_num_cols)),
            remapped_matrix.eval())
  
    def test_load_and_remap_all_missing_rows_and_cols(self):
      """Tests when all the rows & cols are missing and need to be initialized."""
      num_rows = 7
      num_cols = 4
      initializing_values = [42] * num_rows * num_cols
      remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix(
          ckpt_path=[self.bundle_file],
          old_tensor_name=self.old_tensor_name,
          row_remapping=[-1] * num_rows,
          col_remapping=[-1] * num_cols,
          initializing_values=initializing_values,
          num_rows=num_rows,
          num_cols=num_cols)
      with self.test_session():
        self.assertAllClose(
            np.reshape(initializing_values, (num_rows, num_cols)),
            remapped_matrix.eval())
  
    def test_load_and_remap_invalid_remapping(self):
      """Tests that errors are raised when an ID maps to multiple new IDs.
  
      (This should usually not happen when using public APIs).
      """
      invalid_remapping = [1, 0, 0, 0, 1, 2]
  
      # Invalid row remapping.
      remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix(
          ckpt_path=[self.bundle_file],
          old_tensor_name=self.old_tensor_name,
          row_remapping=invalid_remapping,
          col_remapping=[],
          initializing_values=[],
          num_rows=len(invalid_remapping),
          num_cols=self.old_num_cols)
      with self.test_session(), self.assertRaises(errors.UnimplementedError):
        remapped_matrix.eval()
  
      # Invalid column remapping.
      remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix(
          ckpt_path=[self.bundle_file],
          old_tensor_name=self.old_tensor_name,
          row_remapping=list(range(self.old_num_rows)),
          col_remapping=invalid_remapping,
          initializing_values=[],
          num_rows=self.old_num_rows,
          num_cols=len(invalid_remapping))
      with self.test_session(), self.assertRaises(errors.UnimplementedError):
        remapped_matrix.eval()
  
    def test_load_and_remap_incorrect_initializing_values(self):
      """Tests that errors are raised with incorrect number of init values."""
      remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix(
          ckpt_path=[self.bundle_file],
          old_tensor_name=self.old_tensor_name,
          row_remapping=[2, -1, 0],
          col_remapping=[1, -1],
          # Too few initializing values - there should be 4. For some reason,
          # initializing_values must contain no element (instead of 3 or fewer) to
          # ensure that a seg fault would reliably occur if the check raising the
          # InvalidArgumentError were not present.
          initializing_values=[],
          num_rows=3,
          num_cols=2)
      with self.test_session(), self.assertRaises(errors.InvalidArgumentError):
        remapped_matrix.eval()
  
      remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix(
          ckpt_path=[self.bundle_file],
          old_tensor_name=self.old_tensor_name,
          row_remapping=[2, -1, 0],
          col_remapping=[1, -1],
          # Too many initializing values - there should be 4.
          initializing_values=[0] * 5,
          num_rows=3,
          num_cols=2)
      with self.test_session(), self.assertRaises(errors.InvalidArgumentError):
        remapped_matrix.eval()
  
  
  class LoadAndRemapMatrixWithMaxRowsTest(test.TestCase):
    """Tests for the load_and_remap_matrix() op.
  
    (Specifically focused on the max_rows_in_memory arg and its effects on
    TensorBundle's BundleReader and TensorSlice logic).
    """
  
    def _test_loading_variable_with_max_rows(self, np_value, partitioner,
                                             max_rows_in_memory):
      """Helper function for various tests using max_rows_in_memory."""
      ops.reset_default_graph()
      old_tensor_name = 'matrix_to_load_and_remap'
      matrix = variable_scope.get_variable(
          old_tensor_name,
          dtype=dtypes.float32,
          initializer=constant_op.constant(np_value, dtype=dtypes.float32),
          partitioner=partitioner)
  
      with self.test_session() as sess:
        ckpt_path = os.path.join(test.get_temp_dir(), 'temp_ckpt')
        save = saver.Saver([matrix])
        variables.global_variables_initializer().run()
        save.save(sess, ckpt_path)
        num_rows, num_cols = np_value.shape
  
        # Tests loading the entire tensor (except reversed).
        remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix(
            ckpt_path=ckpt_path,
            old_tensor_name=old_tensor_name,
            # Simply reverses the rows of the matrix.
            row_remapping=list(range(num_rows - 1, -1, -1)),
            col_remapping=[],
            initializing_values=[],
            num_rows=num_rows,
            num_cols=num_cols,
            max_rows_in_memory=max_rows_in_memory)
        self.assertAllClose(np_value[::-1], remapped_matrix.eval())
  
        # Tests loading the tensor (except for the first and last rows), with
        # uninitialized values. Requires num_rows to be at least 3 since we're
        # skipping the first and last rows.
        self.assertGreater(num_rows, 2)
        prefix_rows = 2
        suffix_rows = 3
        remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix(
            ckpt_path=ckpt_path,
            old_tensor_name=old_tensor_name,
            # Reverses the rows of the matrix, then prepends and appends
            # uninitialized rows.
            row_remapping=([-1] * prefix_rows + list(range(1, num_rows - 1)) +
                           [-1] * suffix_rows),
            col_remapping=[],
            initializing_values=[42] * (prefix_rows + suffix_rows) * num_cols,
            num_rows=num_rows - 2 + prefix_rows + suffix_rows,
            num_cols=num_cols,
            max_rows_in_memory=max_rows_in_memory)
        self.assertAllClose(
            np.vstack([
                np.tile(42, [prefix_rows, num_cols]), np_value[1:-1],
                np.tile(42, [suffix_rows, num_cols])
            ]), remapped_matrix.eval())
  
        # Tests when everything is taken from initializing_values.
        new_rows = 7
        initializing_values = [42] * new_rows * num_cols
        remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix(
            ckpt_path=ckpt_path,
            old_tensor_name=old_tensor_name,
            # Nothing is loaded from the old tensor.
            row_remapping=[-1] * new_rows,
            col_remapping=[],
            initializing_values=initializing_values,
            num_rows=new_rows,
            num_cols=num_cols,
            max_rows_in_memory=max_rows_in_memory)
        self.assertAllClose(
            np.reshape(initializing_values, (new_rows, num_cols)),
            remapped_matrix.eval())
  
    def test_loading_rows_divisible_by_max_rows(self):
      """Tests loading normal var when rows are evenly divisible by max_rows."""
      self._test_loading_variable_with_max_rows(
          np_value=np.reshape(list(range(0, 36)), (9, 4)),
          partitioner=None,
          # 9 is evenly divisible by 3.
          max_rows_in_memory=3)
  
    def test_loading_rows_not_divisible_by_max_rows(self):
      """Tests loading normal var when rows aren't divisible by max_rows."""
      self._test_loading_variable_with_max_rows(
          np_value=np.reshape(list(range(0, 36)), (9, 4)),
          partitioner=None,
          # 9 is not evenly divisible by 4.
          max_rows_in_memory=4)
  
    def test_loading_rows_less_than_max_rows(self):
      """Tests loading normal var as a single slice.
  
      (When the specified max_rows_in_memory is larger than the number of rows)
      """
      self._test_loading_variable_with_max_rows(
          np_value=np.reshape(list(range(0, 36)), (9, 4)),
          partitioner=None,
          # 10 > 9.
          max_rows_in_memory=10)
  
    def test_loading_no_max_rows(self):
      """Tests loading normal var as a single slice with no valid max_rows."""
      self._test_loading_variable_with_max_rows(
          np_value=np.reshape(list(range(0, 18)), (6, 3)),
          partitioner=None,
          max_rows_in_memory=-1)
  
    def test_loading_partitions_equals_max_rows(self):
      """Tests loading partitioned var sliced on partition boundary."""
      self._test_loading_variable_with_max_rows(
          np_value=np.reshape(list(range(0, 36)), (9, 4)),
          partitioner=partitioned_variables.fixed_size_partitioner(3),
          # With a tensor of shape [9, 3] and 3 partitions, each partition has
          # exactly 3 rows.
          max_rows_in_memory=3)
  
    def test_loading_partitions_greater_than_max_rows(self):
      """Tests loading partitioned var with more slices than partitions."""
      self._test_loading_variable_with_max_rows(
          np_value=np.reshape(list(range(0, 36)), (9, 4)),
          partitioner=partitioned_variables.fixed_size_partitioner(3),
          # Even though each partition has 3 rows, we'll only load the tensor one
          # row at a time.
          max_rows_in_memory=1)
  
    def test_loading_partitions_less_than_max_rows(self):
      """Tests loading partitioned var as a single slice.
  
      (When the specified max_rows_in_memory is larger than the number of rows)
      """
      self._test_loading_variable_with_max_rows(
          np_value=np.reshape(list(range(0, 36)), (9, 4)),
          partitioner=partitioned_variables.fixed_size_partitioner(3),
          max_rows_in_memory=10)
  
    def test_loading_partitions_no_max_rows(self):
      """Tests loading partitioned var as single slice with no valid max_rows."""
      self._test_loading_variable_with_max_rows(
          np_value=np.reshape(list(range(0, 36)), (9, 4)),
          partitioner=partitioned_variables.fixed_size_partitioner(3),
          max_rows_in_memory=-1)
  
  
  if __name__ == '__main__':
    test.main()
  