# Copyright 2015 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 tf.layers.normalization."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf

import os

import numpy as np

from tensorflow.core.protobuf import saver_pb2
from tensorflow.python.framework import test_util
from keras.legacy_tf_layers import convolutional as conv_layers
from keras.legacy_tf_layers import normalization as normalization_layers


@test_util.run_v1_only('b/120545219')
class BNTest(tf.test.TestCase):

  def _simple_model(self, image, fused, freeze_mode):
    output_channels, kernel_size = 2, 3
    conv = conv_layers.conv2d(
        image,
        output_channels,
        kernel_size,
        use_bias=False,
        kernel_initializer=tf.compat.v1.ones_initializer())
    bn_layer = normalization_layers.BatchNormalization(fused=fused)
    bn_layer._bessels_correction_test_only = False
    training = not freeze_mode
    bn = bn_layer.apply(conv, training=training)
    loss = tf.reduce_sum(tf.abs(bn))
    optimizer = tf.compat.v1.train.GradientDescentOptimizer(0.01)
    if not freeze_mode:
      update_ops = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS)
      with tf.control_dependencies(update_ops):
        train_op = optimizer.minimize(loss)
    else:
      train_op = optimizer.minimize(loss)
    saver = tf.compat.v1.train.Saver(write_version=saver_pb2.SaverDef.V2)
    return loss, train_op, saver

  def _train(self,
             checkpoint_path,
             shape,
             use_gpu,
             is_fused,
             restore=False,
             freeze_mode=False,
             dtype=tf.float32):
    tf.compat.v1.reset_default_graph()
    graph = tf.compat.v1.get_default_graph()
    with self.session(graph=graph, use_gpu=use_gpu) as sess:
      image = tf.compat.v1.placeholder(dtype=dtype, shape=shape)
      loss, train_op, saver = self._simple_model(image, is_fused, freeze_mode)
      if restore:
        saver.restore(sess, checkpoint_path)
      else:
        self.evaluate(tf.compat.v1.global_variables_initializer())
      np.random.seed(0)
      for _ in range(2):
        image_val = np.random.rand(*shape).astype(dtype.as_numpy_dtype)
        sess.run([loss, train_op], feed_dict={image: image_val})
      if restore:
        all_vars = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.GLOBAL_VARIABLES)
        all_vars_values = [var.eval() for var in all_vars]
        return all_vars_values
      else:
        saver.save(sess, checkpoint_path)

  def _infer(self, checkpoint_path, image_val, shape, use_gpu, is_fused):
    dtype = image_val.dtype
    tf.compat.v1.reset_default_graph()
    graph = tf.compat.v1.get_default_graph()
    with self.session(graph=graph, use_gpu=use_gpu) as sess:
      image = tf.compat.v1.placeholder(dtype=dtype, shape=shape)
      loss, _, saver = self._simple_model(image, is_fused, True)
      saver.restore(sess, checkpoint_path)
      loss_val = sess.run(loss, feed_dict={image: image_val})
      return loss_val

  def _trainEvalSequence(self, dtype, train1_use_gpu, train2_use_gpu,
                         infer_use_gpu):
    batch, height, width, input_channels = 2, 4, 5, 3
    shape = [batch, height, width, input_channels]

    # Not all characters in a dtype string representation are allowed in
    # filenames in all operating systems. This map will sanitize these.
    dtype_to_valid_fn = {
        tf.float16: 'float16',
        tf.float32: 'float32',
    }
    checkpoint = os.path.join(
        self.get_temp_dir(), 'cp_%s_%s_%s_%s' % (
            dtype_to_valid_fn[dtype], train1_use_gpu, train2_use_gpu,
            infer_use_gpu))

    self._train(
        checkpoint,
        shape,
        use_gpu=train1_use_gpu,
        is_fused=True,
        restore=False,
        freeze_mode=False,
        dtype=dtype)

    train_vars = self._train(
        checkpoint,
        shape,
        use_gpu=train2_use_gpu,
        is_fused=True,
        restore=True,
        freeze_mode=False,
        dtype=dtype)

    np.random.seed(0)
    image_val = np.random.rand(batch, height, width, input_channels).astype(
        dtype.as_numpy_dtype)
    loss_val = self._infer(
        checkpoint, image_val, shape, use_gpu=infer_use_gpu, is_fused=True)

    return train_vars, loss_val

  def testHalfPrecision(self):
    ref_vars, ref_loss = self._trainEvalSequence(
        dtype=tf.float32,
        train1_use_gpu=True,
        train2_use_gpu=True,
        infer_use_gpu=True)

    self.assertEqual(len(ref_vars), 5)

    for train1_use_gpu in [True, False]:
      for train2_use_gpu in [True, False]:
        for infer_use_gpu in [True, False]:
          test_vars, test_loss = self._trainEvalSequence(
              tf.float16, train1_use_gpu, train2_use_gpu, infer_use_gpu)
          self.assertEqual(len(test_vars), 5)
          for test_var, ref_var in zip(test_vars, ref_vars):
            self.assertAllClose(test_var, ref_var, rtol=1.e-3, atol=1.e-3)
          self.assertAllClose(test_loss, ref_loss, rtol=1.e-3, atol=1.e-3)

  def _testCheckpoint(self, is_fused_checkpoint_a, is_fused_checkpoint_b,
                      use_gpu_checkpoint_a, use_gpu_checkpoint_b,
                      use_gpu_test_a, use_gpu_test_b, freeze_mode):
    batch, height, width, input_channels = 2, 4, 5, 3
    shape = [batch, height, width, input_channels]
    base_path = '%s_%s_%s_%s_%s_%s' % (is_fused_checkpoint_a,
                                       is_fused_checkpoint_b,
                                       use_gpu_checkpoint_a,
                                       use_gpu_checkpoint_b, use_gpu_test_a,
                                       use_gpu_test_b)

    checkpoint_path_a = os.path.join(self.get_temp_dir(),
                                     'checkpoint_a_%s' % base_path)
    self._train(
        checkpoint_path_a,
        shape,
        use_gpu_checkpoint_a,
        is_fused_checkpoint_a,
        restore=False,
        freeze_mode=freeze_mode)
    checkpoint_path_b = os.path.join(self.get_temp_dir(),
                                     'checkpoint_b_%s' % base_path)
    self._train(
        checkpoint_path_b,
        shape,
        use_gpu_checkpoint_b,
        is_fused_checkpoint_b,
        restore=False,
        freeze_mode=freeze_mode)

    vars_fused = self._train(
        checkpoint_path_a,
        shape,
        use_gpu_test_a,
        True,
        restore=True,
        freeze_mode=freeze_mode)
    vars_nonfused = self._train(
        checkpoint_path_b,
        shape,
        use_gpu_test_b,
        False,
        restore=True,
        freeze_mode=freeze_mode)
    self.assertEqual(len(vars_fused), 5)
    self.assertEqual(len(vars_nonfused), 5)
    for var_fused, var_nonfused in zip(vars_fused, vars_nonfused):
      self.assertAllClose(var_fused, var_nonfused, atol=1e-5)

    image_val = np.random.rand(batch, height, width,
                               input_channels).astype(np.float32)
    loss_fused_val = self._infer(checkpoint_path_a, image_val, shape,
                                 use_gpu_test_a, True)
    loss_nonfused_val = self._infer(checkpoint_path_b, image_val, shape,
                                    use_gpu_test_b, False)
    self.assertAllClose(loss_fused_val, loss_nonfused_val, atol=1e-6, rtol=3e-4)

  def _testCheckpointCrossDevice(self, ckpt_a_fused, ckpt_a_use_gpu,
                                 ckpt_b_fused, ckpt_b_use_gpu):
    for use_gpu_test_a in [True, False]:
      for use_gpu_test_b in [True, False]:
        for freeze_mode in [True, False]:
          self._testCheckpoint(ckpt_a_fused, ckpt_a_use_gpu, ckpt_b_fused,
                               ckpt_b_use_gpu, use_gpu_test_a, use_gpu_test_b,
                               freeze_mode)

  def testCheckpointFusedCPUAndFusedGPU(self):
    self._testCheckpointCrossDevice(True, False, True, True)

  def testCheckpointFusedCPUAndFusedCPU(self):
    self._testCheckpointCrossDevice(True, False, True, False)

  def testCheckpointFusedGPUAndFusedGPU(self):
    self._testCheckpointCrossDevice(True, True, True, True)

  def testCheckpointNonFusedCPUAndNonFusedGPU(self):
    self._testCheckpointCrossDevice(False, False, False, True)

  def testCheckpointNonFusedCPUAndNonFusedCPU(self):
    self._testCheckpointCrossDevice(False, False, False, False)

  def testCheckpointNonFusedGPUAndNonFusedGPU(self):
    self._testCheckpointCrossDevice(False, True, False, True)

  def testCheckpointNonFusedGPUAndFusedGPU(self):
    self._testCheckpointCrossDevice(False, True, True, True)

  def testCheckpointNonFusedGPUAndFusedCPU(self):
    self._testCheckpointCrossDevice(False, True, True, False)

  def testCheckpointNonFusedCPUAndFusedCPU(self):
    self._testCheckpointCrossDevice(False, False, True, False)

  def testCreateBN(self):
    # Call layer.
    bn = normalization_layers.BatchNormalization(axis=1)
    inputs = tf.random.uniform((5, 4, 3), seed=1)
    training = tf.compat.v1.placeholder(dtype='bool')
    outputs = bn.apply(inputs, training=training)

    # Verify shape.
    self.assertListEqual(outputs.get_shape().as_list(), [5, 4, 3])

    # Verify layer attributes.
    self.assertEqual(len(bn.updates), 2)
    self.assertEqual(len(bn.variables), 4)
    self.assertEqual(len(bn.trainable_variables), 2)
    self.assertEqual(len(bn.non_trainable_variables), 2)

    # Test that updates were created and added to UPDATE_OPS.
    self.assertEqual(len(bn.updates), 2)
    self.assertListEqual(
        tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS), bn.updates)

    # Test that weights were created and added to TRAINABLE_VARIABLES.
    self.assertListEqual(
        tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.TRAINABLE_VARIABLES),
        bn.trainable_variables)

  def testCreateFusedBNFloat16(self):
    # Call layer.
    bn = normalization_layers.BatchNormalization(axis=1, fused=True)
    inputs = tf.random.uniform(
        (5, 4, 3, 3), seed=1, dtype=tf.float16)
    training = tf.compat.v1.placeholder(dtype='bool')
    outputs = bn.apply(inputs, training=training)

    # Verify shape.
    self.assertListEqual(outputs.get_shape().as_list(), [5, 4, 3, 3])

    # Verify layer attributes.
    self.assertEqual(len(bn.updates), 2)
    self.assertEqual(len(bn.variables), 4)
    self.assertEqual(len(bn.trainable_variables), 2)
    self.assertEqual(len(bn.non_trainable_variables), 2)
    for var in bn.variables:
      self.assertTrue(var.dtype._is_ref_dtype)

    # Test that updates were created and added to UPDATE_OPS.
    self.assertEqual(len(bn.updates), 2)
    self.assertListEqual(
        tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS), bn.updates)

    # Test that weights were created and added to TRAINABLE_VARIABLES.
    self.assertListEqual(
        tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.TRAINABLE_VARIABLES),
        bn.trainable_variables)

  def test3DInputAxis1(self):
    epsilon = 1e-3
    bn = normalization_layers.BatchNormalization(
        axis=1, epsilon=epsilon, momentum=0.9)
    inputs = tf.Variable(
        np.random.random((5, 4, 3)) + 100, dtype=tf.float32)
    training = tf.compat.v1.placeholder(dtype='bool')
    outputs = bn.apply(inputs, training=training)

    with self.cached_session() as sess:
      # Test training with placeholder learning phase.
      self.evaluate(tf.compat.v1.global_variables_initializer())

      np_gamma, np_beta = self.evaluate([bn.gamma, bn.beta])
      np_gamma = np.reshape(np_gamma, (1, 4, 1))
      np_beta = np.reshape(np_beta, (1, 4, 1))

      for _ in range(100):
        np_output, _, _ = sess.run([outputs] + bn.updates,
                                   feed_dict={training: True})
        # Verify that the axis is normalized during training.
        normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
        self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
        self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

      # Verify that the statistics are updated during training.
      moving_mean, moving_var = self.evaluate(
          [bn.moving_mean, bn.moving_variance])
      np_inputs = self.evaluate(inputs)
      mean = np.mean(np_inputs, axis=(0, 2))
      std = np.std(np_inputs, axis=(0, 2))
      variance = np.square(std)
      self.assertAllClose(mean, moving_mean, atol=1e-2)
      self.assertAllClose(variance, moving_var, atol=1e-2)

      # Test inference with placeholder learning phase.
      np_output = sess.run(outputs, feed_dict={training: False})

      # Verify that the axis is normalized during inference.
      normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
      self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
      self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

  def test3DInputAxis2(self):
    epsilon = 1e-3
    bn = normalization_layers.BatchNormalization(
        axis=2, epsilon=epsilon, momentum=0.9)
    inputs = tf.Variable(
        np.random.random((5, 4, 3)) + 100, dtype=tf.float32)
    training = tf.compat.v1.placeholder(dtype='bool')
    outputs = bn.apply(inputs, training=training)

    with self.cached_session() as sess:
      # Test training with placeholder learning phase.
      self.evaluate(tf.compat.v1.global_variables_initializer())
      np_gamma, np_beta = self.evaluate([bn.gamma, bn.beta])
      np_gamma = np.reshape(np_gamma, (1, 1, 3))
      np_beta = np.reshape(np_beta, (1, 1, 3))
      for _ in range(100):
        np_output, _, _ = sess.run([outputs] + bn.updates,
                                   feed_dict={training: True})
        # Verify that the axis is normalized during training.
        normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
        self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
        self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

      # Verify that the statistics are updated during training.
      moving_mean, moving_var = self.evaluate(
          [bn.moving_mean, bn.moving_variance])
      np_inputs = self.evaluate(inputs)
      mean = np.mean(np_inputs, axis=(0, 1))
      std = np.std(np_inputs, axis=(0, 1))
      variance = np.square(std)
      self.assertAllClose(mean, moving_mean, atol=1e-2)
      self.assertAllClose(variance, moving_var, atol=1e-2)

      # Test inference with placeholder learning phase.
      np_output = sess.run(outputs, feed_dict={training: False})

      # Verify that the axis is normalized during inference.
      normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
      self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
      self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

  def test4DInputAxis1(self):
    if tf.test.is_gpu_available(cuda_only=True):
      epsilon = 1e-3
      bn = normalization_layers.BatchNormalization(
          axis=1, epsilon=epsilon, momentum=0.9)
      inputs = tf.Variable(
          np.random.random((5, 4, 3, 6)) + 100, dtype=tf.float32)
      training = tf.compat.v1.placeholder(dtype='bool')
      outputs = bn.apply(inputs, training=training)

      with self.session(use_gpu=True) as sess:
        # Test training with placeholder learning phase.
        self.evaluate(tf.compat.v1.global_variables_initializer())
        np_gamma, np_beta = self.evaluate([bn.gamma, bn.beta])
        np_gamma = np.reshape(np_gamma, (1, 4, 1, 1))
        np_beta = np.reshape(np_beta, (1, 4, 1, 1))
        for _ in range(100):
          np_output, _, _ = sess.run(
              [outputs] + bn.updates, feed_dict={training: True})
          # Verify that the axis is normalized during training.
          normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
          self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
          self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

        # Verify that the statistics are updated during training.
        moving_mean, moving_var = self.evaluate(
            [bn.moving_mean, bn.moving_variance])
        np_inputs = self.evaluate(inputs)
        mean = np.mean(np_inputs, axis=(0, 2, 3))
        std = np.std(np_inputs, axis=(0, 2, 3))
        variance = np.square(std)
        self.assertAllClose(mean, moving_mean, atol=1e-2)
        self.assertAllClose(variance, moving_var, atol=1e-2)

        # Test inference with placeholder learning phase.
        np_output = sess.run(outputs, feed_dict={training: False})

        # Verify that the axis is normalized during inference.
        normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
        self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
        self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

  def test4DInputAxis2(self):
    epsilon = 1e-3
    bn = normalization_layers.BatchNormalization(
        axis=2, epsilon=epsilon, momentum=0.9)
    inputs = tf.Variable(
        np.random.random((5, 4, 3, 6)) + 100, dtype=tf.float32)
    training = tf.compat.v1.placeholder(dtype='bool')
    outputs = bn.apply(inputs, training=training)

    with self.cached_session() as sess:
      # Test training with placeholder learning phase.
      self.evaluate(tf.compat.v1.global_variables_initializer())
      np_gamma, np_beta = self.evaluate([bn.gamma, bn.beta])
      np_gamma = np.reshape(np_gamma, (1, 1, 3, 1))
      np_beta = np.reshape(np_beta, (1, 1, 3, 1))
      for _ in range(100):
        np_output, _, _ = sess.run([outputs] + bn.updates,
                                   feed_dict={training: True})
        # Verify that the axis is normalized during training.
        normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
        self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
        self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

      # Verify that the statistics are updated during training.
      moving_mean, moving_var = self.evaluate(
          [bn.moving_mean, bn.moving_variance])
      np_inputs = self.evaluate(inputs)
      mean = np.mean(np_inputs, axis=(0, 1, 3))
      std = np.std(np_inputs, axis=(0, 1, 3))
      variance = np.square(std)
      self.assertAllClose(mean, moving_mean, atol=1e-2)
      self.assertAllClose(variance, moving_var, atol=1e-2)

      # Test inference with placeholder learning phase.
      np_output = sess.run(outputs, feed_dict={training: False})

      # Verify that the axis is normalized during inference.
      normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
      self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
      self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

  def test4DInputAxis3(self):
    epsilon = 1e-3
    bn = normalization_layers.BatchNormalization(
        axis=3, epsilon=epsilon, momentum=0.9)
    inputs = tf.Variable(
        np.random.random((5, 4, 3, 6)) + 100, dtype=tf.float32)
    training = tf.compat.v1.placeholder(dtype='bool')
    outputs = bn.apply(inputs, training=training)

    with self.cached_session() as sess:
      # Test training with placeholder learning phase.
      self.evaluate(tf.compat.v1.global_variables_initializer())
      np_gamma, np_beta = self.evaluate([bn.gamma, bn.beta])
      np_gamma = np.reshape(np_gamma, (1, 1, 1, 6))
      np_beta = np.reshape(np_beta, (1, 1, 1, 6))
      for _ in range(100):
        np_output, _, _ = sess.run([outputs] + bn.updates,
                                   feed_dict={training: True})
        # Verify that the axis is normalized during training.
        normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
        self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
        self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

      # Verify that the statistics are updated during training.
      moving_mean, moving_var = self.evaluate(
          [bn.moving_mean, bn.moving_variance])
      np_inputs = self.evaluate(inputs)
      mean = np.mean(np_inputs, axis=(0, 1, 2))
      std = np.std(np_inputs, axis=(0, 1, 2))
      variance = np.square(std)
      self.assertAllClose(mean, moving_mean, atol=1e-2)
      self.assertAllClose(variance, moving_var, atol=1e-2)

      # Test inference with placeholder learning phase.
      np_output = sess.run(outputs, feed_dict={training: False})

      # Verify that the axis is normalized during inference.
      normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
      self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
      self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

  def test4DInputAxis3Fused(self):
    epsilon = 1e-3
    bn = normalization_layers.BatchNormalization(
        axis=3, epsilon=epsilon, momentum=0.9, fused=True)
    inputs = tf.Variable(
        np.random.random((5, 4, 3, 6)) + 100, dtype=tf.float32)
    training = tf.compat.v1.placeholder(dtype='bool')
    outputs = bn.apply(inputs, training=training)

    with self.cached_session() as sess:
      # Test training with placeholder learning phase.
      self.evaluate(tf.compat.v1.global_variables_initializer())
      np_gamma, np_beta = self.evaluate([bn.gamma, bn.beta])
      np_gamma = np.reshape(np_gamma, (1, 1, 1, 6))
      np_beta = np.reshape(np_beta, (1, 1, 1, 6))
      for _ in range(100):
        np_output, _, _ = sess.run(
            [outputs] + bn.updates, feed_dict={training: True})
        # Verify that the axis is normalized during training.
        normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
        self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
        self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

      # Verify that the statistics are updated during training.
      moving_mean, moving_var = self.evaluate(
          [bn.moving_mean, bn.moving_variance])
      np_inputs = self.evaluate(inputs)
      mean = np.mean(np_inputs, axis=(0, 1, 2))
      std = np.std(np_inputs, axis=(0, 1, 2))
      variance = np.square(std)
      self.assertAllClose(mean, moving_mean, atol=1e-2)
      self.assertAllClose(variance, moving_var, atol=1e-2)

      # Test inference with placeholder learning phase.
      np_output = sess.run(outputs, feed_dict={training: False})

      # Verify that the axis is normalized during inference.
      normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
      self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
      self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

  def test4DInputAxis1Fused(self):
    if tf.test.is_gpu_available(cuda_only=True):
      epsilon = 1e-3
      bn = normalization_layers.BatchNormalization(
          axis=1, epsilon=epsilon, momentum=0.9, fused=True)
      inputs = tf.Variable(
          np.random.random((5, 4, 3, 6)) + 100, dtype=tf.float32)
      training = tf.compat.v1.placeholder(dtype='bool')
      outputs = bn.apply(inputs, training=training)

      with self.cached_session() as sess:
        # Test training with placeholder learning phase.
        self.evaluate(tf.compat.v1.global_variables_initializer())
        np_gamma, np_beta = self.evaluate([bn.gamma, bn.beta])
        np_gamma = np.reshape(np_gamma, (1, 4, 1, 1))
        np_beta = np.reshape(np_beta, (1, 4, 1, 1))
        for _ in range(100):
          np_output, _, _ = sess.run(
              [outputs] + bn.updates, feed_dict={training: True})
          # Verify that the axis is normalized during training.
          normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
          self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
          self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

        # Verify that the statistics are updated during training.
        moving_mean, moving_var = self.evaluate(
            [bn.moving_mean, bn.moving_variance])
        np_inputs = self.evaluate(inputs)
        mean = np.mean(np_inputs, axis=(0, 2, 3))
        std = np.std(np_inputs, axis=(0, 2, 3))
        variance = np.square(std)
        self.assertAllClose(mean, moving_mean, atol=1e-2)
        self.assertAllClose(variance, moving_var, atol=1e-2)

        # Test inference with placeholder learning phase.
        np_output = sess.run(outputs, feed_dict={training: False})

        # Verify that the axis is normalized during inference.
        normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
        self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
        self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

  def testNegativeAxis(self):
    epsilon = 1e-3
    bn = normalization_layers.BatchNormalization(
        axis=-1, epsilon=epsilon, momentum=0.9)
    inputs = tf.Variable(
        np.random.random((5, 4, 3, 6)) + 100, dtype=tf.float32)
    training = tf.compat.v1.placeholder(dtype='bool')
    outputs = bn.apply(inputs, training=training)

    with self.cached_session() as sess:
      # Test training with placeholder learning phase.
      self.evaluate(tf.compat.v1.global_variables_initializer())
      np_gamma, np_beta = self.evaluate([bn.gamma, bn.beta])
      np_gamma = np.reshape(np_gamma, (1, 1, 1, 6))
      np_beta = np.reshape(np_beta, (1, 1, 1, 6))
      for _ in range(100):
        np_output, _, _ = sess.run([outputs] + bn.updates,
                                   feed_dict={training: True})

        # Verify that the axis is normalized during training.
        normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
        self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
        self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

      # Verify that the statistics are updated during training.
      moving_mean, moving_var = self.evaluate(
          [bn.moving_mean, bn.moving_variance])
      np_inputs = self.evaluate(inputs)
      mean = np.mean(np_inputs, axis=(0, 1, 2))
      std = np.std(np_inputs, axis=(0, 1, 2))
      variance = np.square(std)
      self.assertAllClose(mean, moving_mean, atol=1e-2)
      self.assertAllClose(variance, moving_var, atol=1e-2)

      # Test inference with placeholder learning phase.
      np_output = sess.run(outputs, feed_dict={training: False})

      # Verify that the axis is normalized during inference.
      normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
      self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
      self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

  def testBooleanLearningPhase(self):
    epsilon = 1e-3
    bn = normalization_layers.BatchNormalization(
        axis=-1, epsilon=epsilon, momentum=0.9)
    inputs = tf.Variable(
        np.random.random((5, 4, 3, 6)) + 100, dtype=tf.float32)
    outputs_training = bn.apply(inputs, training=True)
    outputs_infer = bn.apply(inputs, training=False)

    with self.cached_session() as sess:
      # Test training with placeholder learning phase.
      self.evaluate(tf.compat.v1.global_variables_initializer())
      np_gamma, np_beta = self.evaluate([bn.gamma, bn.beta])
      np_gamma = np.reshape(np_gamma, (1, 1, 1, 6))
      np_beta = np.reshape(np_beta, (1, 1, 1, 6))
      for _ in range(100):
        np_output, _, _ = sess.run([outputs_training] + bn.updates)
        # Verify that the axis is normalized during training.
        normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
        self.assertAlmostEqual(np.mean(normed_np_output), 0., places=2)
        self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

      # Verify that the statistics are updated during training.
      moving_mean, moving_var = self.evaluate(
          [bn.moving_mean, bn.moving_variance])
      np_inputs = self.evaluate(inputs)
      mean = np.mean(np_inputs, axis=(0, 1, 2))
      std = np.std(np_inputs, axis=(0, 1, 2))
      variance = np.square(std)
      self.assertAllClose(mean, moving_mean, atol=1e-2)
      self.assertAllClose(variance, moving_var, atol=1e-2)

      # Test inference with placeholder learning phase.
      np_output = self.evaluate(outputs_infer)

      # Verify that the axis is normalized during inference.
      normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
      self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
      self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

  def testFunctionalNoReuse(self):
    inputs = tf.Variable(
        np.random.random((5, 4, 3, 6)), dtype=tf.float32)
    epsilon = 1e-3
    training = tf.compat.v1.placeholder(dtype='bool')
    outputs = normalization_layers.batch_norm(
        inputs,
        axis=-1,
        momentum=0.9,
        epsilon=epsilon,
        training=training,
        name='bn')

    updates = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS)
    all_vars = dict([(v.name, v) for v in tf.compat.v1.global_variables()])
    moving_mean = all_vars['bn/moving_mean:0']
    moving_variance = all_vars['bn/moving_variance:0']
    beta = all_vars['bn/beta:0']
    gamma = all_vars['bn/gamma:0']

    with self.cached_session() as sess:
      # Test training with placeholder learning phase.
      self.evaluate(tf.compat.v1.global_variables_initializer())
      np_gamma, np_beta = self.evaluate([gamma, beta])
      np_gamma = np.reshape(np_gamma, (1, 1, 1, 6))
      np_beta = np.reshape(np_beta, (1, 1, 1, 6))
      for _ in range(100):
        np_output, _, _ = sess.run([outputs] + updates,
                                   feed_dict={training: True})
        # Verify that the axis is normalized during training.
        normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
        self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
        self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

      # Verify that the statistics are updated during training.
      np_moving_mean, np_moving_var = self.evaluate(
          [moving_mean, moving_variance])
      np_inputs = self.evaluate(inputs)
      np_mean = np.mean(np_inputs, axis=(0, 1, 2))
      np_std = np.std(np_inputs, axis=(0, 1, 2))
      np_variance = np.square(np_std)
      self.assertAllClose(np_mean, np_moving_mean, atol=1e-2)
      self.assertAllClose(np_variance, np_moving_var, atol=1e-2)

      # Test inference with placeholder learning phase.
      np_output = sess.run(outputs, feed_dict={training: False})

      # Verify that the axis is normalized during inference.
      normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
      self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
      self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

  def testFunctionalReuse(self):
    inputs1 = tf.Variable(
        np.random.random((5, 4, 3, 6)), dtype=tf.float32)
    inputs2 = tf.Variable(
        np.random.random((5, 4, 3, 6)), dtype=tf.float32)
    epsilon = 1e-3
    training = tf.compat.v1.placeholder(dtype='bool')
    _ = normalization_layers.batch_norm(
        inputs1,
        axis=-1,
        momentum=0.9,
        epsilon=epsilon,
        training=training,
        name='bn')
    outputs2 = normalization_layers.batch_norm(
        inputs2,
        axis=-1,
        momentum=0.9,
        epsilon=epsilon,
        training=training,
        name='bn',
        reuse=True)

    # Last 2 update ops
    updates = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS)[-2:]
    all_vars = dict([(v.name, v) for v in tf.compat.v1.global_variables()])
    moving_mean = all_vars['bn/moving_mean:0']
    moving_variance = all_vars['bn/moving_variance:0']
    beta = all_vars['bn/beta:0']
    gamma = all_vars['bn/gamma:0']

    with self.cached_session() as sess:
      # Test training with placeholder learning phase.
      self.evaluate(tf.compat.v1.global_variables_initializer())
      for _ in range(100):
        np_output, _, _ = sess.run([outputs2] + updates,
                                   feed_dict={training: True})

      # Verify that the statistics are updated during training.
      np_moving_mean, np_moving_var = self.evaluate(
          [moving_mean, moving_variance])
      np_inputs = self.evaluate(inputs2)
      np_mean = np.mean(np_inputs, axis=(0, 1, 2))
      np_std = np.std(np_inputs, axis=(0, 1, 2))
      np_variance = np.square(np_std)
      self.assertAllClose(np_mean, np_moving_mean, atol=1e-2)
      self.assertAllClose(np_variance, np_moving_var, atol=1e-2)

      # Verify that the axis is normalized during training.
      np_gamma, np_beta = self.evaluate([gamma, beta])
      np_gamma = np.reshape(np_gamma, (1, 1, 1, 6))
      np_beta = np.reshape(np_beta, (1, 1, 1, 6))
      normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
      self.assertAlmostEqual(np.mean(normed_np_output), 0., places=2)
      self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

      # Test inference with placeholder learning phase.
      np_output = sess.run(outputs2, feed_dict={training: False})

      # Verify that the axis is normalized during inference.
      normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
      self.assertAlmostEqual(np.mean(normed_np_output), 0., places=2)
      self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

  def testFunctionalReuseFromScope(self):
    inputs = tf.Variable(
        np.random.random((5, 4, 3, 6)), dtype=tf.float32)
    epsilon = 1e-3
    training = tf.compat.v1.placeholder(dtype='bool')
    with tf.compat.v1.variable_scope('scope'):
      _ = normalization_layers.batch_norm(
          inputs, axis=-1, momentum=0.9, epsilon=epsilon, training=training)
      self.assertEqual(len(tf.compat.v1.global_variables()), 5)
    with tf.compat.v1.variable_scope('scope', reuse=True):
      _ = normalization_layers.batch_norm(
          inputs, axis=-1, momentum=0.9, epsilon=epsilon, training=training)
      self.assertEqual(len(tf.compat.v1.global_variables()), 5)

  def testNoCenter(self):
    bn = normalization_layers.BatchNormalization(axis=1, center=False)
    inputs = tf.random.uniform((5, 4, 3), seed=1)
    training = tf.compat.v1.placeholder(dtype='bool')
    outputs = bn.apply(inputs, training=training)

    # Verify shape.
    self.assertListEqual(outputs.get_shape().as_list(), [5, 4, 3])

    # Verify layer attributes.
    self.assertEqual(len(bn.updates), 2)
    self.assertEqual(len(bn.variables), 3)
    self.assertEqual(len(bn.trainable_variables), 1)
    self.assertEqual(len(bn.non_trainable_variables), 2)

  def testNoScale(self):
    bn = normalization_layers.BatchNormalization(axis=1, scale=False)
    inputs = tf.random.uniform((5, 4, 3), seed=1)
    training = tf.compat.v1.placeholder(dtype='bool')
    outputs = bn.apply(inputs, training=training)

    # Verify shape.
    self.assertListEqual(outputs.get_shape().as_list(), [5, 4, 3])

    # Verify layer attributes.
    self.assertEqual(len(bn.updates), 2)
    self.assertEqual(len(bn.variables), 3)
    self.assertEqual(len(bn.trainable_variables), 1)
    self.assertEqual(len(bn.non_trainable_variables), 2)

  def testRegularizers(self):
    reg = lambda x: 0.1 * tf.reduce_sum(x)
    bn = normalization_layers.BatchNormalization(axis=1, beta_regularizer=reg)
    inputs = tf.random.uniform((5, 4, 3), seed=1)
    training = tf.compat.v1.placeholder(dtype='bool')
    _ = bn.apply(inputs, training=training)
    self.assertEqual(len(bn.losses), 1)

    bn = normalization_layers.BatchNormalization(axis=1, gamma_regularizer=reg)
    inputs = tf.random.uniform((5, 4, 3), seed=1)
    training = tf.compat.v1.placeholder(dtype='bool')
    _ = bn.apply(inputs, training=training)
    self.assertEqual(len(bn.losses), 1)

  def testConstraints(self):
    g_constraint = lambda x: x / tf.reduce_sum(x)
    b_constraint = lambda x: x / tf.reduce_max(x)
    bn = normalization_layers.BatchNormalization(axis=1,
                                                 gamma_constraint=g_constraint,
                                                 beta_constraint=b_constraint)
    inputs = tf.random.uniform((5, 4, 3), seed=1)
    bn(inputs)
    self.assertEqual(bn.gamma_constraint, g_constraint)
    self.assertEqual(bn.beta_constraint, b_constraint)

  def testRenorm(self):
    shape = (4, 3)
    xt = tf.compat.v1.placeholder(tf.float32, shape)
    momentum = 0.99
    renorm_momentum = 0.8
    rmax = 1.1
    rmin = 0.9
    dmax = 0.1
    gamma = 2.
    beta = 3.
    epsilon = 0.001
    bn = normalization_layers.BatchNormalization(
        axis=1,
        gamma_initializer=tf.compat.v1.constant_initializer(gamma),
        beta_initializer=tf.compat.v1.constant_initializer(beta),
        epsilon=epsilon,
        momentum=momentum,
        renorm=True,
        renorm_clipping={'rmax': rmax, 'rmin': rmin, 'dmax': dmax},
        renorm_momentum=renorm_momentum)
    training = tf.compat.v1.placeholder(tf.bool)
    yt = bn.apply(xt, training=training)

    moving_mean = 0.
    moving_stddev = 1.
    renorm_mean = 0.
    renorm_stddev = 1.
    with self.session(use_gpu=True) as sess:
      self.evaluate(tf.compat.v1.global_variables_initializer())
      for _ in range(5):
        x = np.random.random(shape)

        mean = x.mean(0)
        variance = x.var(0)
        stddev = np.sqrt(variance + epsilon)
        r = (stddev / renorm_stddev).clip(rmin, rmax)
        d = ((mean - renorm_mean) / renorm_stddev).clip(-dmax, dmax)
        y_train = ((x - mean) / stddev * r + d) * gamma + beta
        renorm_mean += (mean - renorm_mean) * (1. - renorm_momentum)
        renorm_stddev += (stddev - renorm_stddev) * (1. - renorm_momentum)
        moving_mean += (mean - moving_mean) * (1. - momentum)
        moving_stddev += (stddev - moving_stddev) * (1. - momentum)

        y_test = ((x - moving_mean) /
                  (moving_stddev * moving_stddev)**0.5 * gamma) + beta

        yt_val_train, _, _ = sess.run([yt] + bn.updates,
                                      feed_dict={xt: x, training: True})
        yt_val_test, _, _ = sess.run([yt] + bn.updates,
                                     feed_dict={xt: x, training: False})

        self.assertAllClose(y_train, yt_val_train, atol=1e-5)
        self.assertAllClose(y_test, yt_val_test, atol=1e-5)

  def testRenormNoClippingSameMomentumGivesSameTestTrain(self):
    shape = (4, 3)
    xt = tf.compat.v1.placeholder(tf.float32, shape)
    momentum = 0.9
    renorm_momentum = 0.9
    gamma = 2.
    beta = 3.
    epsilon = 0.001
    bn = normalization_layers.BatchNormalization(
        axis=1,
        gamma_initializer=tf.compat.v1.constant_initializer(gamma),
        beta_initializer=tf.compat.v1.constant_initializer(beta),
        epsilon=epsilon,
        momentum=momentum,
        renorm=True,
        renorm_clipping=None,
        renorm_momentum=momentum)
    training = tf.compat.v1.placeholder(tf.bool)
    yt = bn.apply(xt, training=training)
    moving_mean = 0.
    moving_stddev = 1.
    renorm_mean = 0.
    renorm_stddev = 1.
    with self.session(use_gpu=True) as sess:
      self.evaluate(tf.compat.v1.global_variables_initializer())
      for step in range(6):
        x = np.random.random(shape)

        mean = x.mean(0)
        variance = x.var(0)
        stddev = np.sqrt(variance + epsilon)
        r = (stddev / renorm_stddev)
        d = ((mean - renorm_mean) / renorm_stddev)
        y_test = ((x - moving_mean) /
                  (moving_stddev * moving_stddev)**0.5 * gamma) + beta
        y_train = ((x - mean) / stddev * r + d) * gamma + beta
        renorm_mean += (mean - renorm_mean) * (1. - renorm_momentum)
        renorm_stddev += (stddev - renorm_stddev) * (1. - renorm_momentum)
        moving_mean += (mean - moving_mean) * (1. - momentum)
        moving_stddev += (stddev - moving_stddev) * (1. - momentum)

        # Compute test values first, before the train mode updates the moving
        # averages.
        yt_val_test, _, _ = sess.run([yt] + bn.updates,
                                     feed_dict={xt: x, training: False})
        yt_val_train, _, _ = sess.run([yt] + bn.updates,
                                      feed_dict={xt: x, training: True})

        # Due to initialization inconsistencies, values may not be identical
        # on the first iteration (but shouldn't be different by much more than
        # epsilon). After the first iteration they should be identical.
        atol = epsilon * 1.5 if step == 0 else 1e-5
        self.assertAllClose(y_train, yt_val_train, atol=atol)
        self.assertAllClose(y_test, yt_val_test, atol=atol)
        self.assertAllClose(yt_val_train, yt_val_test, atol=atol)

  def testAdjustment(self):
    shape = (4, 3)
    xt = tf.compat.v1.placeholder(tf.float32, shape)
    momentum = 0.99
    gamma = 2.
    beta = 3.
    epsilon = 0.001
    adjust_scale = tf.random.uniform(shape[-1:], 0.5, 1.5)
    adjust_bias = tf.random.uniform(shape[-1:], -.2, .2)
    bn = normalization_layers.BatchNormalization(
        axis=1,
        gamma_initializer=tf.compat.v1.constant_initializer(gamma),
        beta_initializer=tf.compat.v1.constant_initializer(beta),
        epsilon=epsilon,
        momentum=momentum,
        adjustment=lambda _: (adjust_scale, adjust_bias))
    training = tf.compat.v1.placeholder(tf.bool)
    yt = bn.apply(xt, training=training)

    moving_mean = 0.
    moving_variance = 1.
    with self.session(use_gpu=True) as sess:
      self.evaluate(tf.compat.v1.global_variables_initializer())
      for _ in range(5):
        x = np.random.random(shape)
        yt_val_train, adj_scale_val, adj_bias_val = sess.run(
            [yt, adjust_scale, adjust_bias] + bn.updates,
            feed_dict={xt: x, training: True})[:3]
        yt_val_test = sess.run([yt] + bn.updates,
                               feed_dict={xt: x, training: False})[0]

        mean = x.mean(0)
        variance = x.var(0)
        y_train = (((x - mean) / (variance + epsilon) ** 0.5) * adj_scale_val +
                   adj_bias_val) * gamma + beta
        moving_mean += (mean - moving_mean) * (1. - momentum)
        moving_variance += (variance - moving_variance) * (1. - momentum)

        y_test = ((x - moving_mean) / (moving_variance + epsilon) ** 0.5 *
                  gamma) + beta

        self.assertAllClose(y_train, yt_val_train, atol=1e-5)
        self.assertAllClose(y_test, yt_val_test, atol=1e-5)

  def testRenormWithAdjustment(self):
    shape = (4, 3)
    xt = tf.compat.v1.placeholder(tf.float32, shape)
    momentum = 0.99
    renorm_momentum = 0.8
    rmax = 1.1
    rmin = 0.9
    dmax = 0.1
    gamma = 2.
    beta = 3.
    epsilon = 0.001
    adjust_scale = tf.random.uniform(shape[-1:], 0.5, 1.5)
    adjust_bias = tf.random.uniform(shape[-1:], -.2, .2)
    bn = normalization_layers.BatchNormalization(
        axis=1,
        gamma_initializer=tf.compat.v1.constant_initializer(gamma),
        beta_initializer=tf.compat.v1.constant_initializer(beta),
        epsilon=epsilon,
        momentum=momentum,
        renorm=True,
        renorm_clipping={'rmax': rmax, 'rmin': rmin, 'dmax': dmax},
        renorm_momentum=renorm_momentum,
        adjustment=lambda _: (adjust_scale, adjust_bias))
    training = tf.compat.v1.placeholder(tf.bool)
    yt = bn.apply(xt, training=training)

    moving_mean = 0.
    moving_stddev = 1.
    renorm_mean = 0.
    renorm_stddev = 1.
    with self.session(use_gpu=True) as sess:
      self.evaluate(tf.compat.v1.global_variables_initializer())
      for _ in range(5):
        x = np.random.random(shape)
        yt_val_train, adj_scale_val, adj_bias_val = sess.run(
            [yt, adjust_scale, adjust_bias] + bn.updates,
            feed_dict={xt: x, training: True})[:3]
        yt_val_test = sess.run([yt] + bn.updates,
                               feed_dict={xt: x, training: False})[0]

        mean = x.mean(0)
        variance = x.var(0)
        stddev = np.sqrt(variance + epsilon)
        r = (stddev / renorm_stddev).clip(rmin, rmax)
        d = ((mean - renorm_mean) / renorm_stddev).clip(-dmax, dmax)
        y_train = (((x - mean) / stddev * r + d) * adj_scale_val +
                   adj_bias_val) * gamma + beta
        renorm_mean += (mean - renorm_mean) * (1. - renorm_momentum)
        renorm_stddev += (stddev - renorm_stddev) * (1. - renorm_momentum)
        moving_mean += (mean - moving_mean) * (1. - momentum)
        moving_stddev += (stddev - moving_stddev) * (1. - momentum)

        y_test = ((x - moving_mean) /
                  (moving_stddev * moving_stddev)**0.5 * gamma) + beta

        self.assertAllClose(y_train, yt_val_train, atol=1e-5)
        self.assertAllClose(y_test, yt_val_test, atol=1e-5)

  def testGhostBNNegativeVirtualBatch(self):
    shape = [6, 5, 4, 3]
    inp = tf.random.uniform(shape, seed=1)

    with self.assertRaises(ValueError):
      normalization_layers.batch_normalization(
          inp, virtual_batch_size=-1)

  def testGhostBNVirtualBatchFull(self):
    shape = [6, 5, 4, 3]
    inp = tf.random.uniform(shape, seed=1)
    out1 = normalization_layers.batch_normalization(inp)
    out2 = normalization_layers.batch_normalization(
        inp, virtual_batch_size=6)

    self.assertListEqual(
        out1.shape.as_list(), out2.shape.as_list())

    with self.session(use_gpu=True) as sess:
      self.evaluate(tf.compat.v1.global_variables_initializer())

      x = np.random.random(shape)
      y1, y2 = sess.run([out1, out2], feed_dict={inp: x})

      self.assertAllClose(y1, y2, atol=1e-5)

  def testGhostBNInputOutputShapesMatch(self):
    shape = [6, 4, 3]
    inp = tf.random.uniform(shape, seed=1)
    out = normalization_layers.batch_normalization(
        inp, virtual_batch_size=3)
    self.assertListEqual(out.shape.as_list(), shape)

  def testGhostBNUnknownBatchSize(self):
    np_shape = [10, 5, 4]
    tf_shape = [None, 5, 4]
    inp = tf.compat.v1.placeholder(tf.float32, tf_shape)
    out = normalization_layers.batch_normalization(
        inp, virtual_batch_size=2)

    with self.session(use_gpu=True) as sess:
      self.evaluate(tf.compat.v1.global_variables_initializer())

      x = np.random.random(np_shape)
      y = sess.run(out, feed_dict={inp: x})

      self.assertListEqual(list(y.shape), np_shape)

  def testGhostBN2Dims(self):
    shape = [6, 2]
    virtual_batch_size = 3
    beta = 2.
    gamma = 3.
    momentum = 0.8
    epsilon = 1e-3
    moving_means = np.zeros([2, 2], dtype=np.float32)
    moving_vars = np.ones([2, 2], dtype=np.float32)

    inp = tf.compat.v1.placeholder(tf.float32, shape)
    is_training = tf.compat.v1.placeholder(tf.bool)
    bn = normalization_layers.BatchNormalization(
        momentum=momentum,
        epsilon=epsilon,
        beta_initializer=tf.compat.v1.constant_initializer(beta),
        gamma_initializer=tf.compat.v1.constant_initializer(gamma),
        virtual_batch_size=virtual_batch_size)
    out = bn.apply(inp, training=is_training)
    ghost_shape = ([virtual_batch_size,
                    shape[0] // virtual_batch_size,
                    shape[1]])

    with self.session(use_gpu=True) as sess:
      self.evaluate(tf.compat.v1.global_variables_initializer())
      for _ in range(5):
        x = np.random.random(shape)

        sub_batched = np.reshape(x, ghost_shape)
        means = np.mean(sub_batched, axis=0, keepdims=True)
        variances = np.var(sub_batched, axis=0, keepdims=True)

        avg_means = np.mean(means, axis=1, keepdims=True)
        avg_variances = np.mean(variances, axis=1, keepdims=True)

        moving_means = moving_means * momentum + avg_means * (1. - momentum)
        moving_vars = moving_vars * momentum + avg_variances * (1. - momentum)

        y_train = ((sub_batched - means) /
                   (variances + epsilon) ** 0.5 * gamma) + beta
        y_test = ((sub_batched - moving_means) /
                  (moving_vars + epsilon) ** 0.5 * gamma) + beta

        y_train = np.reshape(y_train, shape)
        y_test = np.reshape(y_test, shape)

        y_val_train, _, _ = sess.run([out] + bn.updates,
                                     feed_dict={inp: x, is_training: True})
        y_val_test = sess.run(out, feed_dict={inp: x, is_training: False})

        self.assertAllClose(y_train, y_val_train, atol=1e-5)
        self.assertAllClose(y_test, y_val_test, atol=1e-5)

  def testGhostBN4DimsAxis3(self):
    shape = [6, 10, 10, 3]
    virtual_batch_size = 2
    beta = 2.
    gamma = 3.
    momentum = 0.8
    epsilon = 1e-3
    moving_means = np.zeros([1, 1, 1, 1, 3], dtype=np.float32)
    moving_vars = np.ones([1, 1, 1, 1, 3], dtype=np.float32)

    inp = tf.compat.v1.placeholder(tf.float32, shape)
    is_training = tf.compat.v1.placeholder(tf.bool)
    bn = normalization_layers.BatchNormalization(
        axis=3,
        momentum=momentum,
        epsilon=epsilon,
        beta_initializer=tf.compat.v1.constant_initializer(beta),
        gamma_initializer=tf.compat.v1.constant_initializer(gamma),
        virtual_batch_size=virtual_batch_size)
    out = bn.apply(inp, training=is_training)
    ghost_shape = ([virtual_batch_size, shape[0] // virtual_batch_size] +
                   shape[1:])

    with self.session(use_gpu=True) as sess:
      self.evaluate(tf.compat.v1.global_variables_initializer())
      for _ in range(5):
        x = np.random.random(shape)

        sub_batched = np.reshape(x, ghost_shape)
        means = np.mean(sub_batched, axis=(0, 2, 3), keepdims=True)
        variances = np.var(sub_batched, axis=(0, 2, 3), keepdims=True)

        avg_means = np.mean(means, axis=1, keepdims=True)
        avg_variances = np.mean(variances, axis=1, keepdims=True)

        moving_means = moving_means * momentum + avg_means * (1. - momentum)
        moving_vars = moving_vars * momentum + avg_variances * (1. - momentum)

        y_train = ((sub_batched - means) /
                   (variances + epsilon) ** 0.5 * gamma) + beta
        y_test = ((sub_batched - moving_means) /
                  (moving_vars + epsilon) ** 0.5 * gamma) + beta

        y_train = np.reshape(y_train, shape)
        y_test = np.reshape(y_test, shape)

        y_val_train, _, _ = sess.run([out] + bn.updates,
                                     feed_dict={inp: x, is_training: True})
        y_val_test = sess.run(out, feed_dict={inp: x, is_training: False})

        self.assertAllClose(y_train, y_val_train, atol=1e-2)
        self.assertAllClose(y_test, y_val_test, atol=1e-2)

  def testGhostBN4DimsAxis1(self):
    shape = [6, 3, 10, 10]
    virtual_batch_size = 2
    beta = 2.
    gamma = 3.
    momentum = 0.8
    epsilon = 1e-3
    moving_means = np.zeros([1, 1, 3, 1, 1], dtype=np.float32)
    moving_vars = np.ones([1, 1, 3, 1, 1], dtype=np.float32)

    inp = tf.compat.v1.placeholder(tf.float32, shape)
    is_training = tf.compat.v1.placeholder(tf.bool)
    bn = normalization_layers.BatchNormalization(
        axis=1,
        momentum=momentum,
        epsilon=epsilon,
        beta_initializer=tf.compat.v1.constant_initializer(beta),
        gamma_initializer=tf.compat.v1.constant_initializer(gamma),
        virtual_batch_size=virtual_batch_size,
        fused=False)      # NCHW is unsupported by CPU fused batch norm
    out = bn.apply(inp, training=is_training)
    ghost_shape = ([virtual_batch_size, shape[0] // virtual_batch_size] +
                   shape[1:])

    with self.session(use_gpu=True) as sess:
      self.evaluate(tf.compat.v1.global_variables_initializer())
      for _ in range(5):
        x = np.random.random(shape)

        sub_batched = np.reshape(x, ghost_shape)
        means = np.mean(sub_batched, axis=(0, 3, 4), keepdims=True)
        variances = np.var(sub_batched, axis=(0, 3, 4), keepdims=True)

        avg_means = np.mean(means, axis=1, keepdims=True)
        avg_variances = np.mean(variances, axis=1, keepdims=True)

        moving_means = moving_means * momentum + avg_means * (1. - momentum)
        moving_vars = moving_vars * momentum + avg_variances * (1. - momentum)

        y_train = ((sub_batched - means) /
                   (variances + epsilon) ** 0.5 * gamma) + beta
        y_test = ((sub_batched - moving_means) /
                  (moving_vars + epsilon) ** 0.5 * gamma) + beta

        y_train = np.reshape(y_train, shape)
        y_test = np.reshape(y_test, shape)

        y_val_train, _, _ = sess.run([out] + bn.updates,
                                     feed_dict={inp: x, is_training: True})
        y_val_test = sess.run(out, feed_dict={inp: x, is_training: False})

        self.assertAllClose(y_train, y_val_train, atol=1e-2)
        self.assertAllClose(y_test, y_val_test, atol=1e-2)

  def testMultiAxisInvalid(self):
    shape = [6, 5, 4, 3]
    inp = tf.random.uniform(shape, seed=1)

    with self.assertRaises(ValueError):
      normalization_layers.batch_normalization(
          inp, axis=[1, 4])    # out of bounds

    with self.assertRaises(ValueError):
      normalization_layers.batch_normalization(
          inp, axis=[-5, 1])   # out of bounds

    with self.assertRaises(ValueError):
      normalization_layers.batch_normalization(
          inp, axis=[1, 2, 1])   # duplicate

  def test3DInputMultiAxis12(self):
    epsilon = 1e-3
    bn = normalization_layers.BatchNormalization(
        axis=[1, 2], epsilon=epsilon, momentum=0.9)
    inputs = tf.Variable(
        np.random.random((5, 4, 3)) + 100, dtype=tf.float32)
    training = tf.compat.v1.placeholder(dtype='bool')
    outputs = bn.apply(inputs, training=training)

    with self.cached_session() as sess:
      # Test training with placeholder learning phase.
      self.evaluate(tf.compat.v1.global_variables_initializer())

      np_gamma, np_beta = self.evaluate([bn.gamma, bn.beta])

      for _ in range(100):
        np_output, _, _ = sess.run([outputs] + bn.updates,
                                   feed_dict={training: True})
        # Verify that the axis is normalized during training.
        normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
        self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
        self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

      # Verify that the statistics are updated during training.
      moving_mean, moving_var = self.evaluate(
          [bn.moving_mean, bn.moving_variance])
      np_inputs = self.evaluate(inputs)
      mean = np.mean(np_inputs, axis=0, keepdims=True)
      std = np.std(np_inputs, axis=0, keepdims=True)
      variance = np.square(std)
      self.assertAllClose(mean, moving_mean, atol=1e-2)
      self.assertAllClose(variance, moving_var, atol=1e-2)

      # Test inference with placeholder learning phase.
      np_output = sess.run(outputs, feed_dict={training: False})

      # Verify that the axis is normalized during inference.
      normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
      self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
      self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

  def test5DInputMultiAxis123(self):
    epsilon = 1e-3
    bn = normalization_layers.BatchNormalization(
        axis=[1, 2, 3], epsilon=epsilon, momentum=0.9)
    inputs = tf.Variable(
        np.random.random((5, 3, 4, 4, 3)) + 100, dtype=tf.float32)
    training = tf.compat.v1.placeholder(dtype='bool')
    outputs = bn.apply(inputs, training=training)

    with self.cached_session() as sess:
      # Test training with placeholder learning phase.
      self.evaluate(tf.compat.v1.global_variables_initializer())

      np_gamma, np_beta = self.evaluate([bn.gamma, bn.beta])

      for _ in range(100):
        np_output, _, _ = sess.run([outputs] + bn.updates,
                                   feed_dict={training: True})
        # Verify that the axis is normalized during training.
        normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
        self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
        self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

      # Verify that the statistics are updated during training.
      moving_mean, moving_var = self.evaluate(
          [bn.moving_mean, bn.moving_variance])
      np_inputs = self.evaluate(inputs)
      mean = np.mean(np_inputs, axis=(0, 4), keepdims=True)
      std = np.std(np_inputs, axis=(0, 4), keepdims=True)
      variance = np.square(std)
      self.assertAllClose(mean, moving_mean, atol=1e-2)
      self.assertAllClose(variance, moving_var, atol=1e-2)

      # Test inference with placeholder learning phase.
      np_output = sess.run(outputs, feed_dict={training: False})

      # Verify that the axis is normalized during inference.
      normed_np_output = ((np_output - epsilon) * np_gamma) + np_beta
      self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
      self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)

  def testGhostBN5DimsMultiAxis14(self):
    shape = [6, 3, 10, 10, 4]
    virtual_batch_size = 3
    beta = 2.
    gamma = 3.
    momentum = 0.8
    epsilon = 1e-3
    moving_means = np.zeros([1, 1, 3, 1, 1, 4], dtype=np.float32)
    moving_vars = np.ones([1, 1, 3, 1, 1, 4], dtype=np.float32)

    inp = tf.compat.v1.placeholder(tf.float32, shape)
    is_training = tf.compat.v1.placeholder(tf.bool)
    bn = normalization_layers.BatchNormalization(
        axis=[1, 4],
        momentum=momentum,
        epsilon=epsilon,
        beta_initializer=tf.compat.v1.constant_initializer(beta),
        gamma_initializer=tf.compat.v1.constant_initializer(gamma),
        virtual_batch_size=virtual_batch_size,
        fused=False)
    out = bn.apply(inp, training=is_training)
    ghost_shape = ([virtual_batch_size, shape[0] // virtual_batch_size] +
                   shape[1:])

    with self.session(use_gpu=True) as sess:
      self.evaluate(tf.compat.v1.global_variables_initializer())
      for _ in range(5):
        x = np.random.random(shape)

        sub_batched = np.reshape(x, ghost_shape)
        means = np.mean(sub_batched, axis=(0, 3, 4), keepdims=True)
        variances = np.var(sub_batched, axis=(0, 3, 4), keepdims=True)

        avg_means = np.mean(means, axis=1, keepdims=True)
        avg_variances = np.mean(variances, axis=1, keepdims=True)

        moving_means = moving_means * momentum + avg_means * (1. - momentum)
        moving_vars = moving_vars * momentum + avg_variances * (1. - momentum)

        y_train = ((sub_batched - means) /
                   (variances + epsilon) ** 0.5 * gamma) + beta
        y_test = ((sub_batched - moving_means) /
                  (moving_vars + epsilon) ** 0.5 * gamma) + beta

        y_train = np.reshape(y_train, shape)
        y_test = np.reshape(y_test, shape)

        y_val_train, _, _ = sess.run([out] + bn.updates,
                                     feed_dict={inp: x, is_training: True})
        y_val_test = sess.run(out, feed_dict={inp: x, is_training: False})

        self.assertAllClose(y_train, y_val_train, atol=1e-2)
        self.assertAllClose(y_test, y_val_test, atol=1e-2)


if __name__ == '__main__':
  tf.test.main()
