text: # Copyright 2019 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 kfac.utils."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  # Dependency imports
  import numpy as np
  import tensorflow.compat.v1 as tf
  from kfac.python.ops import utils
  
  
  class SequenceDictTest(tf.test.TestCase):
  
    def testSequenceDictInit(self):
      seq_dict = utils.SequenceDict()
      self.assertFalse(seq_dict._dict)
  
    def testSequenceDictInitWithIterable(self):
      reg_dict = {'a': 'foo', 'b': 'bar'}
      itr = zip(reg_dict.keys(), reg_dict.values())
      seq_dict = utils.SequenceDict(itr)
      self.assertEqual(reg_dict, seq_dict._dict)
  
    def testGetItemSingleKey(self):
      seq_dict = utils.SequenceDict({'a': 'foo', 'b': 'bar'})
      self.assertEqual('foo', seq_dict['a'])
  
    def testGetItemMultipleKeys(self):
      seq_dict = utils.SequenceDict({'a': 'foo', 'b': 'bar'})
      self.assertEqual(['foo', 'bar'], seq_dict[('a', 'b')])
  
    def testSetItemSingleKey(self):
      seq_dict = utils.SequenceDict()
      seq_dict['a'] = 'foo'
      self.assertEqual([('a', 'foo')], seq_dict.items())
  
    def testSetItemMultipleKeys(self):
      seq_dict = utils.SequenceDict()
      keys = ('a', 'b', 'c')
      values = ('foo', 'bar', 'baz')
      seq_dict[keys] = values
      self.assertItemsEqual(list(zip(keys, values)), seq_dict.items())
  
  
  class SubGraphTest(tf.test.TestCase):
  
    def testBasicGraph(self):
      a = tf.constant([[1., 2.], [3., 4.]])
      b = tf.constant([[5., 6.], [7., 8.]])
      c = a + b
      d = a * b
      sub_graph = utils.SubGraph((c,))
      self.assertTrue(sub_graph.is_member(a))
      self.assertTrue(sub_graph.is_member(b))
      self.assertTrue(sub_graph.is_member(c))
      self.assertFalse(sub_graph.is_member(d))
  
    def testRepeatedAdds(self):
      a = tf.constant([[1., 2.], [3., 4.]])
      b = tf.constant([[5., 6.], [7., 8.]])
      c = a + b + a  # note that a appears twice in this graph
      sub_graph = utils.SubGraph((c,))
      self.assertTrue(sub_graph.is_member(a))
      self.assertTrue(sub_graph.is_member(b))
      self.assertTrue(sub_graph.is_member(c))
  
    def testFilterList(self):
      a = tf.constant([[1., 2.], [3., 4.]])
      b = tf.constant([[5., 6.], [7., 8.]])
      c = a + b
      d = a * b
      sub_graph = utils.SubGraph((c,))
      input_list = [b, d]
      filtered_list = sub_graph.filter_list(input_list)
      self.assertEqual(filtered_list, [b])
  
    def testVariableUses(self):
      with tf.Graph().as_default():
        var = tf.get_variable('var', shape=[10, 10])
        resource_var = tf.get_variable(
            'resource_var', shape=[10, 10], use_resource=True)
        x = tf.zeros([3, 10])
        z0 = tf.matmul(x, var) + tf.matmul(x, var)
        z1 = tf.matmul(x, resource_var)
        sub_graph = utils.SubGraph((z0, z1))
        self.assertEqual(2, sub_graph.variable_uses(var))
        self.assertEqual(1, sub_graph.variable_uses(resource_var))
  
    def testVariableUsesRelayOps(self):
      with tf.Graph().as_default():
        a = tf.get_variable("a", shape=[2, 2])
        b = tf.get_variable("b", shape=[2, 2])
        ai = tf.identity(a)
        c = tf.matmul(ai, b)
        d = tf.matmul(ai, b)
  
        sub_graph = utils.SubGraph((c, d))
        self.assertEqual(2, sub_graph.variable_uses(a))
        self.assertEqual(2, sub_graph.variable_uses(b))
  
  
  class UtilsTest(tf.test.TestCase):
  
    def _fully_connected_layer_params(self):
      weights_part = tf.constant([[1., 2.], [4., 3.]])
      bias_part = tf.constant([1., 2.])
      return (weights_part, bias_part)
  
    def _conv_layer_params(self):
      weights_shape = 2, 2, 3, 4
      biases_shape = weights_shape[-1:]
      weights = tf.constant(np.random.RandomState(0).randn(*weights_shape))
      biases = tf.constant(np.random.RandomState(1).randn(*biases_shape))
      return (weights, biases)
  
    def testFullyConnectedLayerParamsTupleToMat2d(self):
      with tf.Graph().as_default(), self.test_session() as sess:
        tf.set_random_seed(200)
        layer_params = self._fully_connected_layer_params()
        output = utils.layer_params_to_mat2d(layer_params)
        self.assertListEqual([3, 2], output.get_shape().as_list())
        self.assertAllClose(
            sess.run(output), np.array([[1., 2.], [4., 3.], [1., 2.]]))
  
    def testFullyConnectedLayerParamsTensorToMat2d(self):
      with tf.Graph().as_default(), self.test_session() as sess:
        tf.set_random_seed(200)
        layer_params = self._fully_connected_layer_params()
        output = utils.layer_params_to_mat2d(layer_params[0])
        self.assertListEqual([2, 2], output.get_shape().as_list())
        self.assertAllClose(sess.run(output), np.array([[1., 2.], [4., 3.]]))
  
    def testConvLayerParamsTupleToMat2d(self):
      with tf.Graph().as_default():
        tf.set_random_seed(200)
        layer_params = self._conv_layer_params()
        output = utils.layer_params_to_mat2d(layer_params)
        self.assertListEqual([2 * 2 * 3 + 1, 4], output.get_shape().as_list())
  
    def testKron(self):
      with tf.Graph().as_default(), self.test_session() as sess:
        mat1 = np.array([[1., 2.], [3., 4.]])
        mat2 = np.array([[5., 6.], [7., 8.]])
        mat1_tf = tf.constant(mat1)
        mat2_tf = tf.constant(mat2)
        ans_tf = sess.run(utils.kronecker_product(mat1_tf, mat2_tf))
        ans_np = np.kron(mat1, mat2)
        self.assertAllClose(ans_tf, ans_np)
  
    def testMat2dToFullyConnectedLayerParamsTuple(self):
      with tf.Graph().as_default(), self.test_session() as sess:
        tf.set_random_seed(200)
        vector_template = self._fully_connected_layer_params()
        mat2d = tf.constant([[5., 4.], [3., 2.], [1., 0.]])
  
        output = sess.run(utils.mat2d_to_layer_params(vector_template, mat2d))
  
        self.assertIsInstance(output, tuple)
        self.assertEqual(len(output), 2)
        a, b = output
        self.assertAllClose(a, np.array([[5., 4.], [3., 2.]]))
        self.assertAllClose(b, np.array([1., 0.]))
  
    def testMat2dToFullyConnectedLayerParamsTensor(self):
      with tf.Graph().as_default(), self.test_session() as sess:
        tf.set_random_seed(200)
        vector_template = self._fully_connected_layer_params()[0]
        mat2d = tf.constant([[5., 4.], [3., 2.]])
  
        output = sess.run(utils.mat2d_to_layer_params(vector_template, mat2d))
  
        self.assertAllClose(output, np.array([[5., 4.], [3., 2.]]))
  
    def testTensorsToColumn(self):
      with tf.Graph().as_default(), self.test_session() as sess:
        tf.set_random_seed(200)
  
        vector = tf.constant(np.array([[0., 1.], [2., 3.]]))
        output = utils.tensors_to_column(vector)
        self.assertListEqual([4, 1], output.get_shape().as_list())
        self.assertAllClose(sess.run(output), np.array([0., 1., 2., 3.])[:, None])
  
        vector = self._fully_connected_layer_params()
        output = utils.tensors_to_column(vector)
        self.assertListEqual([6, 1], output.get_shape().as_list())
        self.assertAllClose(
            sess.run(output), np.array([1., 2., 4., 3., 1., 2.])[:, None])
  
        vector = list(vector)
        vector.append(tf.constant([[6.], [7.], [8.], [9.]]))
  
        output = utils.tensors_to_column(vector)
        self.assertListEqual([10, 1], output.get_shape().as_list())
        self.assertAllClose(
            sess.run(output),
            np.array([1., 2., 4., 3., 1., 2., 6., 7., 8., 9.])[:, None])
  
    def testColumnToTensors(self):
      with tf.Graph().as_default(), self.test_session() as sess:
        tf.set_random_seed(200)
  
        vector_template = tf.constant(np.array([[0., 1.], [2., 3.]]))
        colvec = tf.constant(np.arange(4.)[:, None])
        output = sess.run(utils.column_to_tensors(vector_template, colvec))
        self.assertAllClose(output, np.array([[0., 1.], [2., 3.]]))
  
        vector_template = self._fully_connected_layer_params()
        colvec = tf.constant(np.arange(6.)[:, None])
        output = sess.run(utils.column_to_tensors(vector_template, colvec))
  
        self.assertIsInstance(output, tuple)
        self.assertEqual(len(output), 2)
        a, b = output
        self.assertAllClose(a, np.array([[0., 1.], [2., 3.]]))
        self.assertAllClose(b, np.array([4., 5.]))
  
        vector_template = list(vector_template)
        vector_template.append(tf.constant([[6.], [7.], [8.], [9.]]))
        colvec = tf.constant(np.arange(10.)[:, None])
        output = sess.run(utils.column_to_tensors(vector_template, colvec))
        self.assertIsInstance(output, tuple)
        self.assertEqual(len(output), 3)
        a, b, c = output
        self.assertAllClose(a, np.array([[0., 1.], [2., 3.]]))
        self.assertAllClose(b, np.array([4., 5.]))
        self.assertAllClose(c, np.array([[6.], [7.], [8.], [9.]]))
  
    def testPosDefInvCholesky(self):
      with tf.Graph().as_default(), self.test_session() as sess:
        tf.set_random_seed(200)
        np.random.seed(0)
        square = lambda x: np.dot(x, x.T)
  
        size = 3
        x = square(np.random.randn(size, size))
        damp = 0.1
        identity = tf.eye(size, dtype=tf.float64)
  
        tf_inv = utils.posdef_inv_cholesky(tf.constant(x), identity, damp)
        np_inv = np.linalg.inv(x + damp * np.eye(size))
        self.assertAllClose(sess.run(tf_inv), np_inv)
  
    def testPosDefInvMatrixInverse(self):
      with tf.Graph().as_default(), self.test_session() as sess:
        tf.set_random_seed(200)
        np.random.seed(0)
        square = lambda x: np.dot(x, x.T)
  
        size = 3
        x = square(np.random.randn(size, size))
        damp = 0.1
        identity = tf.eye(size, dtype=tf.float64)
  
        tf_inv = utils.posdef_inv_matrix_inverse(tf.constant(x), identity, damp)
        np_inv = np.linalg.inv(x + damp * np.eye(size))
        self.assertAllClose(sess.run(tf_inv), np_inv)
  
    def testBatchExecute(self):
      """Ensure batch_execute runs in a round-robin fashion."""
  
      def increment_var(var):
        return lambda: var.assign_add(1)
  
      with tf.Graph().as_default(), self.test_session() as sess:
        i = tf.get_variable('i', initializer=0)
        accumulators = [
            tf.get_variable('var%d' % j, initializer=0) for j in range(3)
        ]
        thunks = [increment_var(var) for var in accumulators]
        increment_accumulators = utils.batch_execute(i, thunks, 2)
        increment_i = i.assign_add(1)
  
        sess.run(tf.global_variables_initializer())
  
        # Ensure one op per thunk.
        self.assertEqual(3, len(increment_accumulators))
  
        # Ensure round-robin execution.
        values = []
        for _ in range(5):
          sess.run(increment_accumulators)
          sess.run(increment_i)
          values.append(sess.run(accumulators))
        self.assertAllClose(
            [
                [1, 1, 0],  #
                [2, 1, 1],  #
                [2, 2, 2],  #
                [3, 3, 2],  #
                [4, 3, 3]
            ],
            values)
  
    def testExtractConvolutionPatches(self):
      with tf.Graph().as_default(), self.test_session() as sess:
        batch_size = 10
        image_spatial_shape = [9, 10, 11]
        in_channels = out_channels = 32
        kernel_spatial_shape = [5, 3, 3]
        spatial_strides = [1, 2, 1]
        spatial_dilation = [1, 1, 1]
        padding = 'SAME'
  
        images = tf.random_uniform(
            [batch_size] + image_spatial_shape + [in_channels], seed=0)
        kernel_shape = kernel_spatial_shape + [in_channels, out_channels]
        kernel = tf.random_uniform(kernel_shape, seed=1)
  
        # Ensure shape matches expectation.
        patches = utils.extract_convolution_patches(
            images,
            kernel_shape,
            padding,
            strides=spatial_strides,
            dilation_rate=spatial_dilation)
        result_spatial_shape = (
            patches.shape.as_list()[1:1 + len(image_spatial_shape)])
        self.assertEqual(patches.shape.as_list(),
                         [batch_size] + result_spatial_shape +
                         kernel_spatial_shape + [in_channels])
  
        # Ensure extract...patches() + matmul() and convolution() implementation
        # give the same answer.
        outputs = tf.nn.convolution(
            images,
            kernel,
            padding,
            strides=spatial_strides,
            dilation_rate=spatial_dilation)
  
        patches_flat = tf.reshape(
            patches, [-1, np.prod(kernel_spatial_shape) * in_channels])
        kernel_flat = tf.reshape(kernel, [-1, out_channels])
        outputs_flat = tf.matmul(patches_flat, kernel_flat)
  
        outputs_, outputs_flat_ = sess.run([outputs, outputs_flat])
        self.assertAllClose(outputs_.flatten(), outputs_flat_.flatten())
  
    def testExtractPointwiseConv2dPatches(self):
      with tf.Graph().as_default(), self.test_session() as sess:
        batch_size = 10
        image_height = image_width = 8
        in_channels = out_channels = 3
        kernel_height = kernel_width = 1
        strides = [1, 1, 1, 1]
        padding = 'VALID'
  
        images = tf.random_uniform(
            [batch_size, image_height, image_width, in_channels], seed=0)
        kernel_shape = [kernel_height, kernel_width, in_channels, out_channels]
        kernel = tf.random_uniform(kernel_shape, seed=1)
  
        # Ensure shape matches expectation.
        patches = utils.extract_pointwise_conv2d_patches(images, kernel_shape)
        self.assertEqual(patches.shape.as_list(), [
            batch_size, image_height, image_width, kernel_height, kernel_width,
            in_channels
        ])
  
        # Ensure extract...patches() + matmul() and conv2d() implementation
        # give the same answer.
        outputs = tf.nn.conv2d(images, kernel, strides, padding)
  
        patches_flat = tf.reshape(
            patches, [-1, kernel_height * kernel_width * in_channels])
        kernel_flat = tf.reshape(kernel, [-1, out_channels])
        outputs_flat = tf.matmul(patches_flat, kernel_flat)
  
        outputs_, outputs_flat_ = sess.run([outputs, outputs_flat])
        self.assertAllClose(outputs_.flatten(), outputs_flat_.flatten())
  
  
  class AccumulatorVariableTest(tf.test.TestCase):
  
    def test_assign_to_var(self):
      var_shape = (2, 3)
      acc_var = utils.AccumulatorVariable(
          name='test_acc_var', dtype=tf.float32, shape=var_shape)
      values = [
          3. * tf.ones(shape=var_shape), 7. * tf.ones(shape=var_shape),
          11. * tf.ones(shape=var_shape)
      ]
      acc_ops = []
      accc_ops_after_reset = []
      for value in values:
        acc_ops.append(acc_var.accumulate(value))
  
      for value in values[:2]:
        accc_ops_after_reset.append(acc_var.accumulate(value))
  
      init_op = tf.global_variables_initializer()
      with self.test_session() as sess:
  
        sess.run([init_op])
  
        for acc_op in acc_ops:
          sess.run(acc_op)
  
        acc_var_value = sess.run(acc_var.value)
  
        self.assertAllEqual(acc_var_value, 7.*np.ones(shape=var_shape))
  
        sess.run(acc_var.reset())
  
        for acc_op in accc_ops_after_reset:
          sess.run(acc_op)
  
        acc_var_value = sess.run(acc_var.value)
        self.assertAllEqual(acc_var_value, 5. * np.ones(shape=var_shape))
  
    def test_accumulation(self):
      var_shape = (2, 3)
      acc_var = utils.AccumulatorVariable(
          name='test_acc_var', shape=var_shape, dtype=tf.float32)
      values = [
          2. * tf.ones(shape=var_shape), 4. * tf.ones(shape=var_shape),
          9. * tf.ones(shape=var_shape)
      ]
      acc_ops = []
      accc_ops_after_reset = []
      for value in values:
        acc_ops.append(
            acc_var.accumulate(value))
  
      for value in values[:2]:
        accc_ops_after_reset.append(
            acc_var.accumulate(value))
  
      init_op = tf.global_variables_initializer()
      with self.test_session() as sess:
        sess.run([init_op])
  
        for acc_op in acc_ops:
          sess.run([acc_op])
  
        acc_var_value = sess.run(acc_var.read_value_and_reset())
        self.assertAllEqual(acc_var_value, 5. * np.ones(shape=var_shape))
  
        for acc_op in accc_ops_after_reset:
          sess.run([acc_op])
  
        acc_var_value = sess.run(acc_var.value)
        self.assertAllEqual(acc_var_value, 3. * np.ones(shape=var_shape))
  
  
  if __name__ == '__main__':
    tf.disable_v2_behavior()
    tf.test.main()
  