text: # coding=utf-8
  # Copyright 2019 The Google UDA Team Authors.
  #
  # 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.
  """Contains convenience wrappers for typical Neural Network TensorFlow layers.
  
  Ops that have different behavior during training or eval have an is_training
  parameter.
  
  Copied from AutoAugment:
  https://github.com/tensorflow/models/blob/master/research/autoaugment/
  """
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import numpy as np
  import tensorflow as tf
  
  arg_scope = tf.contrib.framework.arg_scope
  FLAGS = tf.flags.FLAGS
  
  
  def variable(name, shape, dtype, initializer, trainable):
    """Returns a TF variable with the passed in specifications."""
    var = tf.get_variable(
        name,
        shape=shape,
        dtype=dtype,
        initializer=initializer,
        trainable=trainable)
    return var
  
  
  def global_avg_pool(x, scope=None):
    """Average pools away spatial height and width dimension of 4D tensor."""
    assert x.get_shape().ndims == 4
    with tf.name_scope(scope, 'global_avg_pool', [x]):
      kernel_size = (1, int(x.shape[1]), int(x.shape[2]), 1)
      squeeze_dims = (1, 2)
      result = tf.nn.avg_pool(
          x,
          ksize=kernel_size,
          strides=(1, 1, 1, 1),
          padding='VALID',
          data_format='NHWC')
      return tf.squeeze(result, squeeze_dims)
  
  
  def zero_pad(inputs, in_filter, out_filter):
    """Zero pads `input` tensor to have `out_filter` number of filters."""
    outputs = tf.pad(inputs, [[0, 0], [0, 0], [0, 0],
                              [(out_filter - in_filter) // 2,
                               (out_filter - in_filter) // 2]])
    return outputs
  
  
  @tf.contrib.framework.add_arg_scope
  def batch_norm(
      inputs,
      update_stats=True,
      decay=0.999,
      center=True,
      scale=False,
      epsilon=0.001,
      is_training=True,
      reuse=None,
      scope=None,
  ):
    """Small wrapper around tf.contrib.layers.batch_norm."""
    batch_norm_op = tf.layers.batch_normalization(
        inputs,
        axis=-1,
        momentum=decay,
        epsilon=epsilon,
        center=center,
        scale=scale,
        training=is_training,
        fused=True,
        trainable=True,
    )
    return batch_norm_op
  
  
  def stride_arr(stride_h, stride_w):
    return [1, stride_h, stride_w, 1]
  
  
  @tf.contrib.framework.add_arg_scope
  def conv2d(inputs,
             num_filters_out,
             kernel_size,
             stride=1,
             scope=None,
             reuse=None):
    """Adds a 2D convolution.
  
      conv2d creates a variable called 'weights', representing the convolutional
      kernel, that is convolved with the input.
  
      Args:
        inputs: a 4D tensor in NHWC format.
        num_filters_out: the number of output filters.
        kernel_size: an int specifying the kernel height and width size.
        stride: an int specifying the height and width stride.
        scope: Optional scope for variable_scope.
        reuse: whether or not the layer and its variables should be reused.
  
      Returns:
        a tensor that is the result of a convolution being applied to `inputs`.
      """
    with tf.variable_scope(scope, 'Conv', [inputs], reuse=reuse):
      num_filters_in = int(inputs.shape[3])
      weights_shape = [kernel_size, kernel_size, num_filters_in, num_filters_out]
  
      # Initialization
      n = int(weights_shape[0] * weights_shape[1] * weights_shape[3])
      weights_initializer = tf.random_normal_initializer(stddev=np.sqrt(2.0 / n))
  
      weights = variable(
          name='weights',
          shape=weights_shape,
          dtype=tf.float32,
          initializer=weights_initializer,
          trainable=True)
      strides = stride_arr(stride, stride)
      outputs = tf.nn.conv2d(
          inputs, weights, strides, padding='SAME', data_format='NHWC')
      return outputs
  
  
  @tf.contrib.framework.add_arg_scope
  def fc(inputs, num_units_out, scope=None, reuse=None):
    """Creates a fully connected layer applied to `inputs`.
  
      Args:
        inputs: a tensor that the fully connected layer will be applied to. It
          will be reshaped if it is not 2D.
        num_units_out: the number of output units in the layer.
        scope: Optional scope for variable_scope.
        reuse: whether or not the layer and its variables should be reused.
  
      Returns:
         a tensor that is the result of applying a linear matrix to `inputs`.
      """
    if len(inputs.shape) > 2:
      inputs = tf.reshape(inputs, [int(inputs.shape[0]), -1])
  
    with tf.variable_scope(scope, 'FC', [inputs], reuse=reuse):
      num_units_in = inputs.shape[1]
      weights_shape = [num_units_in, num_units_out]
      unif_init_range = 1.0 / (num_units_out)**(0.5)
      weights_initializer = tf.random_uniform_initializer(-unif_init_range,
                                                          unif_init_range)
      weights = variable(
          name='weights',
          shape=weights_shape,
          dtype=tf.float32,
          initializer=weights_initializer,
          trainable=True)
      bias_initializer = tf.constant_initializer(0.0)
      biases = variable(
          name='biases',
          shape=[
              num_units_out,
          ],
          dtype=tf.float32,
          initializer=bias_initializer,
          trainable=True)
      outputs = tf.nn.xw_plus_b(inputs, weights, biases)
      return outputs
  
  
  @tf.contrib.framework.add_arg_scope
  def avg_pool(inputs, kernel_size, stride=2, padding='VALID', scope=None):
    """Wrapper around tf.nn.avg_pool."""
    with tf.name_scope(scope, 'AvgPool', [inputs]):
      kernel = stride_arr(kernel_size, kernel_size)
      strides = stride_arr(stride, stride)
      return tf.nn.avg_pool(
          inputs,
          ksize=kernel,
          strides=strides,
          padding=padding,
          data_format='NHWC')
  