from __future__ import division, print_function, absolute_import

from collections import namedtuple

import numpy as np
import six
import numbers
import tensorflow as tf
from tensorflow.python.framework import ops
from tensorflow.python.ops import math_ops
from tensorflow.python.training import moving_averages
from tensorflow.python.framework import function

from . import initializers as initz
from ..utils import util as helper
rng = np.random.RandomState([2018, 1, 1])
NamedOutputs = namedtuple('NamedOutputs', ['name', 'outputs'])


# pylint: disable=redefined-builtin
def input(shape, name='inputs', outputs_collections=None, **unused):
  """Define input layer.

  Args:
      shape: A `Tensor`, define the input shape
          e.g. for image input [batch_size, height, width, depth]
      name: A optional score/name for this op
      outputs_collections: The collections to which the outputs are added.

  Returns:
      A placeholder for the input
  """
  _check_unused(unused, name)
  with tf.name_scope(name):
    inputs = tf.placeholder(tf.float32, shape=shape, name="input")
  return _collect_named_outputs(outputs_collections, name, inputs)


def register_to_collections(inputs, name=None, outputs_collections=None, **unused):
  """Add item to colelction.

  Args:
      shape: A `Tensor`, define the input shape
          e.g. for image input [batch_size, height, width, depth]
      name: A optional score/name for this op
      outputs_collections: The collections to which the outputs are added.

  Returns:
      A placeholder for the input
  """
  _check_unused(unused, name)
  return _collect_named_outputs(outputs_collections, name, inputs)


def fully_connected(x,
                    n_output,
                    is_training,
                    reuse,
                    trainable=True,
                    w_init=initz.he_normal(),
                    b_init=0.0,
                    w_regularizer=tf.nn.l2_loss,
                    w_normalized=False,
                    name='fc',
                    batch_norm=None,
                    batch_norm_args=None,
                    activation=None,
                    params=None,
                    outputs_collections=None,
                    use_bias=True):
  """Adds a fully connected layer.

      `fully_connected` creates a variable called `weights`, representing a fully
      connected weight matrix, which is multiplied by the `x` to produce a
      `Tensor` of hidden units. If a `batch_norm` is provided (such as
      `batch_norm`), it is then applied. Otherwise, if `batch_norm` is
      None and a `b_init` and `use_bias` is provided then a `biases` variable would be
      created and added the hidden units. Finally, if `activation` is not `None`,
      it is applied to the hidden units as well.
      Note: that if `x` have a rank greater than 2, then `x` is flattened
      prior to the initial matrix multiply by `weights`.

  Args:
      x: A `Tensor` of with at least rank 2 and value for the last dimension,
          i.e. `[batch_size, depth]`, `[None, None, None, channels]`.
      is_training: Bool, training or testing
      n_output: Integer or long, the number of output units in the layer.
      reuse: whether or not the layer and its variables should be reused. To be
          able to reuse the layer scope must be given.
      activation: activation function, set to None to skip it and maintain
          a linear activation.
      batch_norm: normalization function to use. If
         `batch_norm` is `True` then google original implementation is used and
          if another function is provided then it is applied.
          default set to None for no normalizer function
      batch_norm_args: normalization function parameters.
      w_init: An initializer for the weights.
      w_regularizer: Optional regularizer for the weights.
      b_init: An initializer for the biases. If None skip biases.
      outputs_collections: The collections to which the outputs are added.
      trainable: If `True` also add variables to the graph collection
          `GraphKeys.TRAINABLE_VARIABLES` (see tf.Variable).
      name: Optional name or scope for variable_scope/name_scope.
      use_bias: Whether to add bias or not

  Returns:
      The 2-D `Tensor` variable representing the result of the series of operations.
      e.g: 2-D `Tensor` [batch, n_output].

  Raises:
      ValueError: if `x` has rank less than 2 or if its last dimension is not set.
  """
  if not (isinstance(n_output, six.integer_types)):
    raise ValueError('n_output should be int or long, got %s.', n_output)

  input_shape = helper.get_input_shape(x)
  assert len(input_shape) > 1, "Input Tensor shape must be > 1-D"
  if len(x.get_shape()) != 2:
    x = _flatten(x)

  n_input = x.get_shape().as_list()[1]
  with tf.variable_scope(name, reuse=reuse):
    shape = [n_input, n_output] if hasattr(w_init, '__call__') else None
    if params is not None:
      W, b = params
    else:
      W = tf.get_variable(
          name='W',
          shape=shape,
          dtype=tf.float32,
          initializer=w_init,
          regularizer=w_regularizer,
          trainable=trainable)
    if w_normalized:
      W = W / tf.reduce_sum(tf.square(W), 1, keep_dims=True)
    output = tf.matmul(x, W)

    if use_bias:
      if params is None:
        b = tf.get_variable(
            name='b',
            shape=[n_output],
            dtype=tf.float32,
            initializer=tf.constant_initializer(b_init),
            trainable=trainable,
        )

      output = tf.nn.bias_add(value=output, bias=b)

    if batch_norm is not None:
      if isinstance(batch_norm, bool):
        batch_norm = batch_norm_tf
      batch_norm_args = batch_norm_args or {}
      output = batch_norm(
          output, is_training=is_training, reuse=reuse, trainable=trainable, **batch_norm_args)

    if activation:
      if activation in (tf.tanh, tf.sigmoid, tf.nn.relu):
        output = activation(output)
      else:
        output = activation(output, reuse=reuse, trainable=trainable)

    return _collect_named_outputs(outputs_collections, name, output)


def conv2d(x,
           n_output_channels,
           is_training,
           reuse,
           trainable=True,
           filter_size=(3, 3),
           stride=(1, 1),
           padding='SAME',
           w_init=initz.he_normal(),
           b_init=0.0,
           w_regularizer=tf.nn.l2_loss,
           untie_biases=False,
           name='conv2d',
           batch_norm=None,
           batch_norm_args=None,
           activation=None,
           use_bias=True,
           outputs_collections=None):
  """Adds a 2D convolutional layer.

      `convolutional layer` creates a variable called `weights`, representing a conv
      weight matrix, which is multiplied by the `x` to produce a
      `Tensor` of hidden units. If a `batch_norm` is provided (such as
      `batch_norm`), it is then applied. Otherwise, if `batch_norm` is
      None and a `b_init` and `use_bias` is provided then a `biases` variable would be
      created and added the hidden units. Finally, if `activation` is not `None`,
      it is applied to the hidden units as well.
      Note: that if `x` have a rank 4

  Args:
      x: A 4-D `Tensor` of with at least rank 2 and value for the last dimension,
          i.e. `[batch_size, in_height, in_width, depth]`,
      is_training: Bool, training or testing
      n_output: Integer or long, the number of output units in the layer.
      reuse: whether or not the layer and its variables should be reused. To be
          able to reuse the layer scope must be given.
      filter_size: a int or list/tuple of 2 positive integers specifying the spatial
          dimensions of of the filters.
      stride: a int or tuple/list of 2 positive integers specifying the stride at which to
          compute output.
      padding: one of `"VALID"` or `"SAME"`.
      activation: activation function, set to None to skip it and maintain
          a linear activation.
      batch_norm: normalization function to use. If
          `batch_norm` is `True` then google original implementation is used and
          if another function is provided then it is applied.
          default set to None for no normalizer function
      batch_norm_args: normalization function parameters.
      w_init: An initializer for the weights.
      w_regularizer: Optional regularizer for the weights.
      untie_biases: spatial dimensions wise baises
      b_init: An initializer for the biases. If None skip biases.
      outputs_collections: The collections to which the outputs are added.
      trainable: If `True` also add variables to the graph collection
          `GraphKeys.TRAINABLE_VARIABLES` (see tf.Variable).
      name: Optional name or scope for variable_scope/name_scope.
      use_bias: Whether to add bias or not

  Returns:
      The 4-D `Tensor` variable representing the result of the series of operations.
      e.g.: 4-D `Tensor` [batch, new_height, new_width, n_output].

  Raises:
      ValueError: if `x` has rank less than 4 or if its last dimension is not set.
  """
  input_shape = helper.get_input_shape(x)
  assert len(input_shape) == 4, "Input Tensor shape must be 4-D"
  with tf.variable_scope(name, reuse=reuse):
    shape = helper.filter_2d(filter_size,
                             x.get_shape()[-1], n_output_channels) if hasattr(w_init,
                                                                              '__call__') else None
    W = tf.get_variable(
        name='W', shape=shape, initializer=w_init, regularizer=w_regularizer, trainable=trainable)

    output = tf.nn.conv2d(
        input=x, filter=W, strides=helper.stride_2d(stride), padding=helper.kernel_padding(padding))

    if use_bias:
      if untie_biases:
        b = tf.get_variable(
            name='b',
            shape=output.get_shape()[1:],
            initializer=tf.constant_initializer(b_init),
            trainable=trainable,
        )
        output = tf.add(output, b)
      else:
        b = tf.get_variable(
            name='b',
            shape=[n_output_channels],
            initializer=tf.constant_initializer(b_init),
            trainable=trainable,
        )
        output = tf.nn.bias_add(value=output, bias=b)

    if batch_norm is not None:
      if isinstance(batch_norm, bool):
        batch_norm = batch_norm_tf
      batch_norm_args = batch_norm_args or {}
      output = batch_norm(
          output, is_training=is_training, reuse=reuse, trainable=trainable, **batch_norm_args)

    if activation:
      output = activation(output, reuse=reuse, trainable=trainable)

    return _collect_named_outputs(outputs_collections, name, output)


def dilated_conv2d(x,
                   n_output_channels,
                   is_training,
                   reuse,
                   trainable=True,
                   filter_size=(3, 3),
                   dilation=1,
                   stride=1,
                   padding='SAME',
                   w_init=initz.he_normal(),
                   b_init=0.0,
                   w_regularizer=tf.nn.l2_loss,
                   untie_biases=False,
                   name='dilated_conv2d',
                   batch_norm=None,
                   batch_norm_args=None,
                   activation=None,
                   use_bias=True,
                   outputs_collections=None):
  """Adds a 2D dilated convolutional layer.

      also known as convolution with holes or atrous convolution.
      If the rate parameter is equal to one, it performs regular 2-D convolution.
      If the rate parameter
      is greater than one, it performs convolution with holes, sampling the input
      values every rate pixels in the height and width dimensions.
      `convolutional layer` creates a variable called `weights`, representing a conv
      weight matrix, which is multiplied by the `x` to produce a
      `Tensor` of hidden units. If a `batch_norm` is provided (such as
      `batch_norm`), it is then applied. Otherwise, if `batch_norm` is
      None and a `b_init` and `use_bias` is provided then a `biases` variable would be
      created and added the hidden units. Finally, if `activation` is not `None`,
      it is applied to the hidden units as well.
      Note: that if `x` have a rank 4

  Args:
      x: A 4-D `Tensor` of with rank 4 and value for the last dimension,
          i.e. `[batch_size, in_height, in_width, depth]`,
      is_training: Bool, training or testing
      n_output: Integer or long, the number of output units in the layer.
      reuse: whether or not the layer and its variables should be reused. To be
          able to reuse the layer scope must be given.
      filter_size: a int or list/tuple of 2 positive integers specifying the spatial
      dimensions of of the filters.
      dilation:  A positive int32. The stride with which we sample input values across
          the height and width dimensions. Equivalently, the rate by which we upsample the
          filter values by inserting zeros across the height and width dimensions. In the literature,
          the same parameter is sometimes called input stride/rate or dilation.
      padding: one of `"VALID"` or `"SAME"`.
      activation: activation function, set to None to skip it and maintain
          a linear activation.
      batch_norm: normalization function to use. If
          `batch_norm` is `True` then google original implementation is used and
          if another function is provided then it is applied.
          default set to None for no normalizer function
      batch_norm_args: normalization function parameters.
      w_init: An initializer for the weights.
      w_regularizer: Optional regularizer for the weights.
      untie_biases: spatial dimensions wise baises
      b_init: An initializer for the biases. If None skip biases.
      outputs_collections: The collections to which the outputs are added.
      trainable: If `True` also add variables to the graph collection
          `GraphKeys.TRAINABLE_VARIABLES` (see tf.Variable).
      name: Optional name or scope for variable_scope/name_scope.
      use_bias: Whether to add bias or not

  Returns:
      The 4-D `Tensor` variable representing the result of the series of operations.
      e.g.: 4-D `Tensor` [batch, new_height, new_width, n_output].

  Raises:
      ValueError: if x has rank less than 4 or if its last dimension is not set.
  """
  input_shape = helper.get_input_shape(x)
  assert len(input_shape) == 4, "Input Tensor shape must be 4-D"
  with tf.variable_scope(name, reuse=reuse):
    shape = helper.filter_2d(filter_size,
                             x.get_shape()[-1], n_output_channels) if hasattr(w_init,
                                                                              '__call__') else None
    W = tf.get_variable(
        name='W', shape=shape, initializer=w_init, regularizer=w_regularizer, trainable=trainable)

    output = tf.nn.atrous_conv2d(
        value=x, filters=W, rate=dilation, padding=helper.kernel_padding(padding))

    if use_bias:
      if untie_biases:
        b = tf.get_variable(
            name='b',
            shape=output.get_shape()[1:],
            initializer=tf.constant_initializer(b_init),
            trainable=trainable,
        )
        output = tf.add(output, b)
      else:
        b = tf.get_variable(
            name='b',
            shape=[n_output_channels],
            initializer=tf.constant_initializer(b_init),
            trainable=trainable,
        )
        output = tf.nn.bias_add(value=output, bias=b)

    if batch_norm is not None:
      if isinstance(batch_norm, bool):
        batch_norm = batch_norm_tf
      batch_norm_args = batch_norm_args or {}
      output = batch_norm(
          output, is_training=is_training, reuse=reuse, trainable=trainable, **batch_norm_args)

    if activation:
      output = activation(output, reuse=reuse, trainable=trainable)

    return _collect_named_outputs(outputs_collections, name, output)


def separable_conv2d(x,
                     n_output_channels,
                     is_training,
                     reuse,
                     trainable=True,
                     filter_size=(3, 3),
                     stride=(1, 1),
                     depth_multiplier=1,
                     padding='SAME',
                     w_init=initz.he_normal(),
                     b_init=0.0,
                     w_regularizer=tf.nn.l2_loss,
                     untie_biases=False,
                     name='separable_conv2d',
                     batch_norm=None,
                     batch_norm_args=None,
                     activation=None,
                     use_bias=True,
                     outputs_collections=None):
  """Adds a 2D seperable convolutional layer.

      Performs a depthwise convolution that acts separately on channels followed by
      a pointwise convolution that mixes channels. Note that this is separability between
      dimensions [1, 2] and 3, not spatial separability between dimensions 1 and 2.
      `convolutional layer` creates two variable called `depthwise_W` and `pointwise_W`,
      `depthwise_W` is multiplied by `x` to produce depthwise conolution, which is multiplied by
      the `pointwise_W` to produce a output `Tensor`
      If a `batch_norm` is provided (such as
      `batch_norm`), it is then applied. Otherwise, if `batch_norm` is
      None and a `b_init` and `use_bias` is provided then a `biases` variable would be
      created and added the hidden units. Finally, if `activation` is not `None`,
      it is applied to the hidden units as well.
      Note: that if `x` have a rank 4

  Args:
      x: A 4-D `Tensor` of with rank 4 and value for the last dimension,
          i.e. `[batch_size, in_height, in_width, depth]`,
      is_training: Bool, training or testing
      n_output: Integer or long, the number of output units in the layer.
      reuse: whether or not the layer and its variables should be reused. To be
          able to reuse the layer scope must be given.
      filter_size: a int or list/tuple of 2 positive integers specifying the spatial
      stride: a int or tuple/list of 2 positive integers specifying the stride at which to
          compute output.
      dimensions of of the filters.
      depth_multiplier:  A positive int32. the number of depthwise convolution output channels for
          each input channel. The total number of depthwise convolution output
          channels will be equal to `num_filters_in * depth_multiplier
      padding: one of `"VALID"` or `"SAME"`.
      activation: activation function, set to None to skip it and maintain
          a linear activation.
      batch_norm: normalization function to use. If
          `batch_norm` is `True` then google original implementation is used and
          if another function is provided then it is applied.
          default set to None for no normalizer function
      batch_norm_args: normalization function parameters.
      w_init: An initializer for the weights.
      w_regularizer: Optional regularizer for the weights.
      untie_biases: spatial dimensions wise baises
      b_init: An initializer for the biases. If None skip biases.
      outputs_collections: The collections to which the outputs are added.
      trainable: If `True` also add variables to the graph collection
          `GraphKeys.TRAINABLE_VARIABLES` (see tf.Variable).
      name: Optional name or scope for variable_scope/name_scope.
      use_bias: Whether to add bias or not

  Returns:
      The 4-D `Tensor` variable representing the result of the series of operations.
      e.g.: 4-D `Tensor` [batch, new_height, new_width, n_output].

  Raises:
      ValueError: if `x` has rank less than 4 or if its last dimension is not set.
  """
  input_shape = helper.get_input_shape(x)
  assert len(input_shape) == 4, "Input Tensor shape must be 4-D"
  with tf.variable_scope(name, reuse=reuse):
    depthwise_shape = helper.filter_2d(filter_size,
                                       x.get_shape()[-1], depth_multiplier) if hasattr(
                                           w_init, '__call__') else None
    pointwise_shape = [1, 1, x.get_shape()[-1] * depth_multiplier, n_output_channels] if hasattr(
        w_init, '__call__') else None
    depthwise_W = tf.get_variable(
        name='depthwise_W',
        shape=depthwise_shape,
        initializer=w_init,
        regularizer=w_regularizer,
        trainable=trainable)
    pointwise_W = tf.get_variable(
        name='pointwise_W',
        shape=pointwise_shape,
        initializer=w_init,
        regularizer=w_regularizer,
        trainable=trainable)

    output = tf.nn.separable_conv2d(
        input=x,
        depthwise_filter=depthwise_W,
        pointwise_filter=pointwise_W,
        strides=helper.stride_2d(stride),
        padding=helper.kernel_padding(padding))

    if use_bias:
      if untie_biases:
        b = tf.get_variable(
            name='b',
            shape=output.get_shape()[1:],
            initializer=tf.constant_initializer(b_init),
            trainable=trainable,
        )
        output = tf.add(output, b)
      else:
        b = tf.get_variable(
            name='b',
            shape=[n_output_channels],
            initializer=tf.constant_initializer(b_init),
            trainable=trainable,
        )
        output = tf.nn.bias_add(value=output, bias=b)

    if batch_norm is not None:
      if isinstance(batch_norm, bool):
        batch_norm = batch_norm_tf
      batch_norm_args = batch_norm_args or {}
      output = batch_norm(
          output, is_training=is_training, reuse=reuse, trainable=trainable, **batch_norm_args)

    if activation:
      output = activation(output, reuse=reuse, trainable=trainable)

    return _collect_named_outputs(outputs_collections, name, output)


def depthwise_conv2d(x,
                     depth_multiplier,
                     is_training,
                     reuse,
                     trainable=True,
                     filter_size=(3, 3),
                     stride=(1, 1),
                     padding='SAME',
                     w_init=initz.he_normal(),
                     b_init=0.0,
                     w_regularizer=tf.nn.l2_loss,
                     untie_biases=False,
                     name='depthwise_conv2d',
                     batch_norm=None,
                     batch_norm_args=None,
                     activation=None,
                     use_bias=True,
                     outputs_collections=None):
  """Adds a 2D sdepthwise convolutional layer.

      Given an input tensor of shape [batch, in_height, in_width, in_channels] and a filter
      tensor of shape [filter_height, filter_width, in_channels, channel_multiplier] containing
      in_channels convolutional filters of depth 1, depthwise_conv2d applies a different filter
      to each input channel (expanding from 1 channel to channel_multiplier channels for each),
      then concatenates the results together. The output has in_channels * channel_multiplier
      channels.
      If a `batch_norm` is provided (such as
      `batch_norm`), it is then applied. Otherwise, if `batch_norm` is
      None and a `b_init` and `use_bias` is provided then a `biases` variable would be
      created and added the hidden units. Finally, if `activation` is not `None`,
      it is applied to the hidden units as well.
      Note: that if `x` have a rank 4

  Args:
      x: A 4-D `Tensor` of with rank 4 and value for the last dimension,
          i.e. `[batch_size, in_height, in_width, depth]`,
      is_training: Bool, training or testing
      reuse: whether or not the layer and its variables should be reused. To be
          able to reuse the layer scope must be given.
      filter_size: a int or list/tuple of 2 positive integers specifying the spatial
          dimensions of of the filters.
      stride: a int or tuple/list of 2 positive integers specifying the stride at which to
          compute output.
      depth_multiplier:  A positive int32. the number of depthwise convolution output channels for
          each input channel. The total number of depthwise convolution output
          channels will be equal to `num_filters_in * depth_multiplier
      padding: one of `"VALID"` or `"SAME"`.
      activation: activation function, set to None to skip it and maintain
          a linear activation.
      batch_norm: normalization function to use. If
          `batch_norm` is `True` then google original implementation is used and
          if another function is provided then it is applied.
          default set to None for no normalizer function
      batch_norm_args: normalization function parameters.
      w_init: An initializer for the weights.
      w_regularizer: Optional regularizer for the weights.
      untie_biases: spatial dimensions wise baises
      b_init: An initializer for the biases. If None skip biases.
      outputs_collections: The collections to which the outputs are added.
      trainable: If `True` also add variables to the graph collection
          `GraphKeys.TRAINABLE_VARIABLES` (see tf.Variable).
      name: Optional name or scope for variable_scope/name_scope.
      use_bias: Whether to add bias or not

  Returns:
      The tensor variable representing the result of the series of operations.
      e.g.: 4-D `Tensor` [batch, new_height, new_width, n_output].

  Raises:
      ValueError: if `x` has rank less than 4 or if its last dimension is not set.
  """
  input_shape = helper.get_input_shape(x)
  assert len(input_shape) == 4, "Input Tensor shape must be 4-D"
  with tf.variable_scope(name, reuse=reuse):
    shape = helper.filter_2d(filter_size,
                             x.get_shape()[-1], depth_multiplier) if hasattr(w_init,
                                                                             '__call__') else None
    W = tf.get_variable(
        name='W', shape=shape, initializer=w_init, regularizer=w_regularizer, trainable=trainable)

    output = tf.nn.depthwise_conv2d(
        x, W, strides=helper.stride_2d(stride), padding=helper.kernel_padding(padding))

    if use_bias:
      if untie_biases:
        b = tf.get_variable(
            name='b',
            shape=output.get_shape()[1:],
            initializer=tf.constant_initializer(b_init),
            trainable=trainable,
        )
        output = tf.add(output, b)
      else:
        b = tf.get_variable(
            name='b',
            shape=[x.get_shape()[-1] * depth_multiplier],
            initializer=tf.constant_initializer(b_init),
            trainable=trainable,
        )
        output = tf.nn.bias_add(value=output, bias=b)

    if batch_norm is not None:
      if isinstance(batch_norm, bool):
        batch_norm = batch_norm_tf
      batch_norm_args = batch_norm_args or {}
      output = batch_norm(
          output, is_training=is_training, reuse=reuse, trainable=trainable, **batch_norm_args)

    if activation:
      output = activation(output, reuse=reuse, trainable=trainable)

    return _collect_named_outputs(outputs_collections, name, output)


def conv3d(x,
           n_output_channels,
           is_training,
           reuse,
           trainable=True,
           filter_size=(3, 3, 3),
           stride=(1, 1, 1),
           padding='SAME',
           w_init=initz.he_normal(),
           b_init=0.0,
           w_regularizer=tf.nn.l2_loss,
           untie_biases=False,
           name='conv3d',
           batch_norm=None,
           batch_norm_args=None,
           activation=None,
           use_bias=True,
           outputs_collections=None):
  """Adds a 3D convolutional layer.

      `convolutional layer` creates a variable called `weights`, representing a conv
      weight matrix, which is multiplied by the `x` to produce a
      `Tensor` of hidden units. If a `batch_norm` is provided (such as
      `batch_norm`), it is then applied. Otherwise, if `batch_norm` is
      None and a `b_init` and `use_bias` is provided then a `biases` variable would be
      created and added the hidden units. Finally, if `activation` is not `None`,
      it is applied to the hidden units as well.
      Note: that if `x` have a rank 5

  Args:
      x: A 5-D `Tensor` of with at least rank 2 and value for the last dimension,
          i.e. `[batch_size, in_depth, in_height, in_width, depth]`,
      is_training: Bool, training or testing
      n_output: Integer or long, the number of output units in the layer.
      reuse: whether or not the layer and its variables should be reused. To be
          able to reuse the layer scope must be given.
      filter_size: a int, or  list/tuple of 3 positive integers specifying the spatial
          dimensions of of the filters.
      stride: a int, or tuple/list of 3 positive integers specifying the stride at which to
          compute output.
      padding: one of `"VALID"` or `"SAME"`.
      activation: activation function, set to None to skip it and maintain
          a linear activation.
      batch_norm: normalization function to use. If
          `batch_norm` is `True` then google original implementation is used and
          if another function is provided then it is applied.
          default set to None for no normalizer function
      batch_norm_args: normalization function parameters.
      w_init: An initializer for the weights.
      w_regularizer: Optional regularizer for the weights.
      untie_biases: spatial dimensions wise baises
      b_init: An initializer for the biases. If None skip biases.
      outputs_collections: The collections to which the outputs are added.
      trainable: If `True` also add variables to the graph collection
          `GraphKeys.TRAINABLE_VARIABLES` (see tf.Variable).
      name: Optional name or scope for variable_scope/name_scope.
      use_bias: Whether to add bias or not

  Returns:
      The 5-D `Tensor` variable representing the result of the series of operations.
      e.g.: 5-D `Tensor` [batch, new_depth, new_height, new_width, n_output].

  Raises:
      ValueError: if `x` has rank less than 5 or if its last dimension is not set.
  """
  input_shape = helper.get_input_shape(x)
  assert len(input_shape) == 5, "Input Tensor shape must be 5-D"
  with tf.variable_scope(name, reuse=reuse):
    shape = helper.filter_3d(filter_size,
                             x.get_shape()[-1], n_output_channels) if hasattr(w_init,
                                                                              '__call__') else None
    W = tf.get_variable(
        name='W', shape=shape, initializer=w_init, regularizer=w_regularizer, trainable=trainable)

    output = tf.nn.conv3d(
        input=x, filter=W, strides=helper.stride_3d(stride), padding=helper.kernel_padding(padding))

    if use_bias:
      if untie_biases:
        b = tf.get_variable(
            name='b',
            shape=output.get_shape()[1:],
            initializer=tf.constant_initializer(b_init),
            trainable=trainable,
        )
        output = tf.add(output, b)
      else:
        b = tf.get_variable(
            name='b',
            shape=[n_output_channels],
            initializer=tf.constant_initializer(b_init),
            trainable=trainable,
        )
        output = tf.nn.bias_add(value=output, bias=b)

    if batch_norm is not None:
      if isinstance(batch_norm, bool):
        batch_norm = batch_norm_tf
      batch_norm_args = batch_norm_args or {}
      output = batch_norm(
          output, is_training=is_training, reuse=reuse, trainable=trainable, **batch_norm_args)

    if activation:
      output = activation(output, reuse=reuse, trainable=trainable)

    return _collect_named_outputs(outputs_collections, name, output)


def upsample2d(input_,
               output_shape,
               is_training,
               reuse,
               trainable=True,
               filter_size=(5, 5),
               stride=(2, 2),
               w_init=initz.he_normal(seed=None),
               b_init=0.0,
               w_regularizer=tf.nn.l2_loss,
               batch_norm=None,
               batch_norm_args=None,
               activation=None,
               name="deconv2d",
               use_bias=True,
               with_w=False,
               outputs_collections=None,
               **unused):
  """Adds a 2D upsampling or deconvolutional layer.

      his operation is sometimes called "deconvolution" after Deconvolutional Networks,
      but is actually the transpose (gradient) of conv2d rather than an actual deconvolution.
      If a `batch_norm` is provided (such as
      `batch_norm`), it is then applied. Otherwise, if `batch_norm` is
      None and a `b_init` and `use_bias` is provided then a `biases` variable would be
      created and added the hidden units. Finally, if `activation` is not `None`,
      it is applied to the hidden units as well.
      Note: that if `x` have a rank 4

  Args:
      x: A 4-D `Tensor` of with at least rank 2 and value for the last dimension,
          i.e. `[batch_size, in_height, in_width, depth]`,
      is_training: Bool, training or testing
      output_shape: 4D tensor, the output shape
      reuse: whether or not the layer and its variables should be reused. To be
          able to reuse the layer scope must be given.
      filter_size: a int or list/tuple of 2 positive integers specifying the spatial
          dimensions of of the filters.
      stride: a int or tuple/list of 2 positive integers specifying the stride at which to
          compute output.
      padding: one of `"VALID"` or `"SAME"`.
      activation: activation function, set to None to skip it and maintain
          a linear activation.
      batch_norm: normalization function to use. If
          `batch_norm` is `True` then google original implementation is used and
          if another function is provided then it is applied.
          default set to None for no normalizer function
      batch_norm_args: normalization function parameters.
      w_init: An initializer for the weights.
      w_regularizer: Optional regularizer for the weights.
      b_init: An initializer for the biases. If None skip biases.
      outputs_collections: The collections to which the outputs are added.
      trainable: If `True` also add variables to the graph collection
          `GraphKeys.TRAINABLE_VARIABLES` (see tf.Variable).
      name: Optional name or scope for variable_scope/name_scope.
      use_bias: Whether to add bias or not

  Returns:
      The tensor variable representing the result of the series of operations.
      e.g.: 4-D `Tensor` [batch, new_height, new_width, n_output].

  Raises:
      ValueError: if `x` has rank less than 4 or if its last dimension is not set.
  """
  input_shape = helper.get_input_shape(input_)
  assert len(input_shape) == 4, "Input Tensor shape must be 4-D"
  with tf.variable_scope(name or 'upsample2d', reuse=reuse):
    shape = helper.filter_2d(filter_size, output_shape[-1],
                             input_.get_shape()[-1]) if hasattr(w_init, '__call__') else None

    # filter : [height, width, output_channels, in_channels]
    w = tf.get_variable(
        name='W', shape=shape, initializer=w_init, regularizer=w_regularizer, trainable=trainable)

    output = tf.nn.conv2d_transpose(
        input_, w, output_shape=output_shape, strides=helper.stride_2d(stride))
    if use_bias:
      biases = tf.get_variable(
          name='biases',
          shape=[output_shape[-1]],
          initializer=tf.constant_initializer(b_init),
          trainable=trainable)
      output = tf.reshape(tf.nn.bias_add(output, biases), output.get_shape())

    if batch_norm is not None:
      if isinstance(batch_norm, bool):
        batch_norm = batch_norm_tf
      batch_norm_args = batch_norm_args or {}
      output = batch_norm(
          output, is_training=is_training, reuse=reuse, trainable=trainable, **batch_norm_args)

    if activation:
      output = activation(output, reuse=reuse)

    tf.add_to_collection(tf.GraphKeys.ACTIVATIONS, output)

    if with_w:
      return _collect_named_outputs(outputs_collections, name, output), w, biases
    else:
      return _collect_named_outputs(outputs_collections, name, output)


def upsample3d(input_,
               output_shape,
               is_training,
               reuse,
               trainable=True,
               filter_size=(5, 5, 5),
               stride=(2, 2, 2),
               w_init=initz.he_normal(seed=None),
               b_init=0.0,
               w_regularizer=tf.nn.l2_loss,
               batch_norm=None,
               batch_norm_args=None,
               activation=None,
               name="deconv3d",
               use_bias=True,
               with_w=False,
               outputs_collections=None,
               **unused):
  """Adds a 3D upsampling or deconvolutional layer.

      his operation is sometimes called "deconvolution" after Deconvolutional Networks,
      but is actually the transpose (gradient) of conv2d rather than an actual deconvolution.
      If a `batch_norm` is provided (such as
      `batch_norm`), it is then applied. Otherwise, if `batch_norm` is
      None and a `b_init` and `use_bias` is provided then a `biases` variable would be
      created and added the hidden units. Finally, if `activation` is not `None`,
      it is applied to the hidden units as well.
      Note: that if `x` have a rank 5

  Args:
      x: A 5-D `Tensor` of with at least rank 2 and value for the last dimension,
          i.e. `[batch_size, in_depth, in_height, in_width, depth]`,
      is_training: Bool, training or testing
      output_shape: 5D tensor, the output shape
      reuse: whether or not the layer and its variables should be reused. To be
          able to reuse the layer scope must be given.
      filter_size: a int or list/tuple of 3 positive integers specifying the spatial
          dimensions of of the filters.
      stride: a int or tuple/list of 3 positive integers specifying the stride at which to
          compute output.
      padding: one of `"VALID"` or `"SAME"`.
      activation: activation function, set to None to skip it and maintain
          a linear activation.
      batch_norm: normalization function to use. If
          `batch_norm` is `True` then google original implementation is used and
          if another function is provided then it is applied.
          default set to None for no normalizer function
      batch_norm_args: normalization function parameters.
      w_init: An initializer for the weights.
      w_regularizer: Optional regularizer for the weights.
      b_init: An initializer for the biases. If None skip biases.
      outputs_collections: The collections to which the outputs are added.
      trainable: If `True` also add variables to the graph collection
          `GraphKeys.TRAINABLE_VARIABLES` (see tf.Variable).
      name: Optional name or scope for variable_scope/name_scope.
      use_bias: Whether to add bias or not

  Returns:
      The tensor variable representing the result of the series of operations.
      e.g.: 5-D `Tensor` [batch, new_depth, new_height, new_width, n_output].

  Raises:
      ValueError: if `x` has rank less than 4 or if its last dimension is not set.
  """
  input_shape = helper.get_input_shape(input_)
  assert len(input_shape) == 5, "Input Tensor shape must be 5-D"
  with tf.variable_scope(name or 'upsample3d', reuse=reuse):
    shape = helper.filter_3d(filter_size, output_shape[-1],
                             input_.get_shape()[-1]) if hasattr(w_init, '__call__') else None

    # filter : [depth, height, width, output_channels, in_channels]
    w = tf.get_variable(
        name='W', shape=shape, initializer=w_init, regularizer=w_regularizer, trainable=trainable)

    output = tf.nn.conv3d_transpose(
        input_, w, output_shape=output_shape, strides=helper.stride_3d(stride))
    if use_bias:
      biases = tf.get_variable(
          name='biases',
          shape=[output_shape[-1]],
          initializer=tf.constant_initializer(b_init),
          trainable=trainable)
      output = tf.reshape(tf.nn.bias_add(output, biases), output.get_shape())

    if batch_norm is not None:
      if isinstance(batch_norm, bool):
        batch_norm = batch_norm_tf
      batch_norm_args = batch_norm_args or {}
      output = batch_norm(
          output, is_training=is_training, reuse=reuse, trainable=trainable, **batch_norm_args)

    if activation:
      output = activation(output, reuse=reuse)

    tf.add_to_collection(tf.GraphKeys.ACTIVATIONS, output)

    if with_w:
      return _collect_named_outputs(outputs_collections, name, output), w, biases
    else:
      return _collect_named_outputs(outputs_collections, name, output)


def conv1d(x,
           n_output_channels,
           is_training,
           reuse,
           trainable=True,
           filter_size=3,
           stride=1,
           padding='SAME',
           w_init=initz.he_normal(),
           b_init=0.0,
           w_regularizer=tf.nn.l2_loss,
           untie_biases=False,
           name='conv1d',
           batch_norm=None,
           batch_norm_args=None,
           activation=None,
           use_bias=True,
           outputs_collections=None):
  """Adds a 1D convolutional layer.

      `convolutional layer` creates a variable called `weights`, representing a conv
      weight matrix, which is multiplied by the `x` to produce a
      `Tensor` of hidden units. If a `batch_norm` is provided (such as
      `batch_norm`), it is then applied. Otherwise, if `batch_norm` is
      None and a `b_init` and `use_bias` is provided then a `biases` variable would be
      created and added the hidden units. Finally, if `activation` is not `None`,
      it is applied to the hidden units as well.
      Note: that if `x` have a rank 4

  Args:
      x: A 3-D `Tensor` of with at least rank 2 and value for the last dimension,
          i.e. `[batch_size, in_width, depth]`,
      is_training: Bool, training or testing
      n_output: Integer or long, the number of output units in the layer.
      reuse: whether or not the layer and its variables should be reused. To be
          able to reuse the layer scope must be given.
      filter_size: a `int specifying the spatial
          dimensions of of the filters.
      stride: a `int` specifying the stride at which to
          compute output.
      padding: one of `"VALID"` or `"SAME"`.
      activation: activation function, set to None to skip it and maintain
          a linear activation.
      batch_norm: normalization function to use. If
          `batch_norm` is `True` then google original implementation is used and
          if another function is provided then it is applied.
          default set to None for no normalizer function
      batch_norm_args: normalization function parameters.
      w_init: An initializer for the weights.
      w_regularizer: Optional regularizer for the weights.
      untie_biases: spatial dimensions wise baises
      b_init: An initializer for the biases. If None skip biases.
      outputs_collections: The collections to which the outputs are added.
      trainable: If `True` also add variables to the graph collection
          `GraphKeys.TRAINABLE_VARIABLES` (see tf.Variable).
      name: Optional name or scope for variable_scope/name_scope.
      use_bias: Whether to add bias or not

  Returns:
      The 3-D `Tensor` variable representing the result of the series of operations.
      e.g.: 3-D `Tensor` [batch, new_width, n_output].

  Raises:
      ValueError: if `x` has rank less than 4 or if its last dimension is not set.
  """
  input_shape = helper.get_input_shape(x)
  assert len(input_shape) == 3, "Input Tensor shape must be 4-D"
  with tf.variable_scope(name, reuse=reuse):
    shape = (filter_size, x.get_shape()[-1],
             n_output_channels) if hasattr(w_init, '__call__') else None
    W = tf.get_variable(
        name='W', shape=shape, initializer=w_init, regularizer=w_regularizer, trainable=trainable)

    output = tf.nn.conv1d(x, W, stride, padding=helper.kernel_padding(padding))

    if use_bias:
      if untie_biases:
        b = tf.get_variable(
            name='b',
            shape=output.get_shape()[1:],
            initializer=tf.constant_initializer(b_init),
            trainable=trainable,
        )
        output = tf.add(output, b)
      else:
        b = tf.get_variable(
            name='b',
            shape=[n_output_channels],
            initializer=tf.constant_initializer(b_init),
            trainable=trainable,
        )
        output = tf.nn.bias_add(value=output, bias=b)

    if batch_norm is not None:
      if isinstance(batch_norm, bool):
        batch_norm = batch_norm_tf
      batch_norm_args = batch_norm_args or {}
      output = batch_norm(
          output, is_training=is_training, reuse=reuse, trainable=trainable, **batch_norm_args)

    if activation:
      output = activation(output, reuse=reuse, trainable=trainable)

    return _collect_named_outputs(outputs_collections, name, output)


def max_pool_1d(x,
                filter_size=3,
                stride=2,
                padding='SAME',
                name='maxpool1d',
                outputs_collections=None,
                **unused):
  """Max Pooling 1D.

  Args:
      x: a 3-D `Tensor` [batch_size, steps, in_channels].
      kernel_size: `int` or `list of int`. Pooling kernel size.
      strides: `int` or `list of int`. Strides of conv operation.
          Default: same as kernel_size.
      padding: `str` from `"same", "valid"`. Padding algo to use.
          Default: 'same'.
      name: A name for this layer (optional). Default: 'maxpool1d'.

  Returns:
      3-D Tensor [batch, pooled steps, in_channels].
  """
  input_shape = helper.get_input_shape(x)
  assert len(input_shape) == 3, "Input Tensor shape must be 3-D"

  filter_size = helper.kernel_2d(filter_size)
  filter_size = [1, filter_size[1], 1, 1]
  stride = helper.stride_2d(stride)
  stride = [1, stride[1], 1, 1]

  with tf.name_scope(name):
    x = tf.expand_dims(x, 2)
    output = tf.nn.max_pool(
        value=x,
        ksize=filter_size,
        strides=stride,
        padding=helper.kernel_padding(padding),
    )
    output = tf.squeeze(output, [2])
    return _collect_named_outputs(outputs_collections, name, output)


def avg_pool_1d(x,
                filter_size=3,
                stride=2,
                padding='SAME',
                name='avgpool1d',
                outputs_collections=None,
                **unused):
  """Avg Pooling 1D.

  Args:
      x: a 3-D `Tensor` [batch_size, steps, in_channels].
      kernel_size: `int` or `list of int`. Pooling kernel size.
      strides: `int` or `list of int`. Strides of conv operation.
          Default: same as kernel_size.
      padding: `str` from `"same", "valid"`. Padding algo to use.
          Default: 'same'.
      name: A name for this layer (optional). Default: 'avgpool1d'.

  Returns:
      3-D Tensor [batch, pooled steps, in_channels].
  """
  input_shape = helper.get_input_shape(x)
  assert len(input_shape) == 3, "Input Tensor shape must be 3-D"

  filter_size = helper.kernel_2d(filter_size)
  filter_size = [1, filter_size[1], 1, 1]
  stride = helper.stride_2d(stride)
  stride = [1, stride[1], 1, 1]

  with tf.name_scope(name):
    x = tf.expand_dims(x, 2)
    output = tf.nn.avg_pool(
        value=x,
        ksize=filter_size,
        strides=stride,
        padding=helper.kernel_padding(padding),
    )
    output = tf.squeeze(output, [2])
    return _collect_named_outputs(outputs_collections, name, output)


def _phase_shift(input_, r):
  bsize, a, b, c = helper.get_input_shape(input_)
  X = tf.reshape(input_, (bsize, a, b, r, r))
  X = tf.transpose(X, (0, 1, 2, 4, 3))
  X = tf.split(1, a, X)
  X = tf.concat([tf.squeeze(x) for x in X], axis=2)
  X = tf.split(1, b, X)
  X = tf.concat([tf.squeeze(x) for x in X], axis=2)
  output = tf.reshape(X, (bsize, a * r, b * r, 1))
  return output


def subpixel2d(input_, r, color=True, name=None, outputs_collections=None, **unused):
  input_shape = helper.get_input_shape(input_)
  assert len(input_shape) == 4, "Input Tensor shape must be 4-D"
  with tf.name_scope(name or "subpixel"):
    if color:
      inputc = tf.split(3, 3, input_)
      output = tf.concat([_phase_shift(x, r) for x in inputc], axis=3)
    else:
      output = _phase_shift(input_, r)
  return _collect_named_outputs(outputs_collections, name, output)


def highway_conv2d(x,
                   n_output,
                   is_training,
                   reuse,
                   trainable=True,
                   filter_size=(3, 3),
                   stride=(1, 1),
                   padding='SAME',
                   w_init=initz.he_normal(),
                   b_init=0.0,
                   w_regularizer=tf.nn.l2_loss,
                   name='highway_conv2d',
                   batch_norm=None,
                   batch_norm_args=None,
                   activation=tf.nn.relu,
                   use_bias=True,
                   outputs_collections=None):
  """Adds a 2D highway convolutional layer.

      https://arxiv.org/abs/1505.00387
      If a `batch_norm` is provided (such as
      `batch_norm`), it is then applied. Otherwise, if `batch_norm` is
      None and a `b_init` and `use_bias` is provided then a `biases` variable would be
      created and added the hidden units. Finally, if `activation` is not `None`,
      it is applied to the hidden units as well.
      Note: that if `x` have a rank 4

  Args:
      x: A 4-D `Tensor` of with at least rank 2 and value for the last dimension,
          i.e. `[batch_size, in_height, in_width, depth]`,
      is_training: Bool, training or testing
      n_output: Integer or long, the number of output units in the layer.
      reuse: whether or not the layer and its variables should be reused. To be
          able to reuse the layer scope must be given.
      filter_size: a int or list/ tuple of 2 positive integers specifying the spatial
          dimensions of of the filters.
      stride: a int or tuple/list of 2 positive integers specifying the stride at which to
          compute output.
      padding: one of `"VALID"` or `"SAME"`.
      activation: activation function, set to None to skip it and maintain
          a linear activation.
      batch_norm: normalization function to use. If
          `batch_norm` is `True` then google original implementation is used and
          if another function is provided then it is applied.
          default set to None for no normalizer function
      batch_norm_args: normalization function parameters.
      w_init: An initializer for the weights.
      w_regularizer: Optional regularizer for the weights.
      untie_biases: spatial dimensions wise baises
      b_init: An initializer for the biases. If None skip biases.
      outputs_collections: The collections to which the outputs are added.
      trainable: If `True` also add variables to the graph collection
          `GraphKeys.TRAINABLE_VARIABLES` (see tf.Variable).
      name: Optional name or scope for variable_scope/name_scope.
      use_bias: Whether to add bias or not

  Returns:
      The `Tensor` variable representing the result of the series of operations.
      e.g.: 4-D `Tensor` [batch, new_height, new_width, n_output].

  Raises:
      ValueError: if `x` has rank less than 4 or if its last dimension is not set.
  """
  input_shape = helper.get_input_shape(x)
  assert len(input_shape) == 4, "Input Tensor shape must be 4-D"
  with tf.variable_scope(name, reuse=reuse):
    w_shape = helper.filter_2d(filter_size,
                               x.get_shape()[-1], n_output) if hasattr(w_init, '__call__') else None

    w_t_shape = [n_output]
    b_shape = [n_output]
    W, b = helper.weight_bias(
        w_shape,
        b_shape,
        w_init=w_init,
        b_init=b_init,
        w_regularizer=w_regularizer,
        trainable=trainable,
        name='main_gate/')
    W_t, b_t = helper.weight_bias(
        w_t_shape,
        b_shape,
        w_init=w_init,
        b_init=-3.0,
        w_regularizer=None,
        trainable=trainable,
        name='transform_gate/')
    output_conv2d = tf.nn.conv2d(
        input=x, filter=W, strides=helper.stride_2d(stride), padding=helper.kernel_padding(padding))
    H = activation(output_conv2d + b, reuse=reuse, trainable=trainable, name='activation_H')
    T = tf.sigmoid(tf.multiply(output_conv2d, W_t) + b_t, name='transform_gate')
    C = tf.subtract(1.0, T, name="carry_gate")
    output = tf.add(tf.multiply(H, T), tf.multiply(output_conv2d, C), name='output')
    if batch_norm is not None:
      if isinstance(batch_norm, bool):
        batch_norm = batch_norm_tf
      batch_norm_args = batch_norm_args or {}
      output = batch_norm(
          output, is_training=is_training, reuse=reuse, trainable=trainable, **batch_norm_args)

    if activation:
      output = activation(output, reuse=reuse, trainable=trainable)

    return _collect_named_outputs(outputs_collections, name, output)


def highway_fc2d(x,
                 n_output,
                 is_training,
                 reuse,
                 trainable=True,
                 filter_size=(3, 3),
                 w_init=initz.he_normal(),
                 b_init=0.0,
                 w_regularizer=tf.nn.l2_loss,
                 name='highway_fc2d',
                 activation=None,
                 use_bias=True,
                 outputs_collections=None):
  """Adds a fully connected highway layer.

      https://arxiv.org/abs/1505.00387
      If a `batch_norm` is provided (such as
      `batch_norm`), it is then applied. Otherwise, if `batch_norm` is
      None and a `b_init` and `use_bias` is provided then a `biases` variable would be
      created and added the hidden units. Finally, if `activation` is not `None`,
      it is applied to the hidden units as well.
      Note: that if `x` have a rank greater than 2, then `x` is flattened
      prior to the initial matrix multiply by `weights`.

  Args:
      x: A 2-D/4-D `Tensor` of with at least rank 2 and value for the last dimension,
          i.e. `[batch_size, depth]`, `[None, None, None, channels]`.
      is_training: Bool, training or testing
      n_output: Integer or long, the number of output units in the layer.
      reuse: whether or not the layer and its variables should be reused. To be
          able to reuse the layer scope must be given.
      activation: activation function, set to None to skip it and maintain
          a linear activation.
      batch_norm: normalization function to use. If
          `batch_norm` is `True` then google original implementation is used and
          if another function is provided then it is applied.
          default set to None for no normalizer function
      batch_norm_args: normalization function parameters.
      w_init: An initializer for the weights.
      w_regularizer: Optional regularizer for the weights.
      b_init: An initializer for the biases. If None skip biases.
      outputs_collections: The collections to which the outputs are added.
      trainable: If `True` also add variables to the graph collection
          `GraphKeys.TRAINABLE_VARIABLES` (see tf.Variable).
      name: Optional name or scope for variable_scope/name_scope.
      use_bias: Whether to add bias or not

  Returns:
      The 2-D `Tensor` variable representing the result of the series of operations.
      e.g.: 2-D `Tensor` [batch_size, n_output]

  Raises:
      ValueError: if `x` has rank less than 2 or if its last dimension is not set.
  """
  input_shape = helper.get_input_shape(x)
  assert len(input_shape) > 1, "Input Tensor shape must be > 1-D"
  if len(x.get_shape()) != 2:
    x = _flatten(x)

  n_input = x.get_shape().as_list()[1]
  w_shape = [n_input, n_output] if hasattr(w_init, '__call__') else None
  b_shape = [n_output]
  with tf.variable_scope(name, reuse=reuse):
    W, b = helper.weight_bias(
        w_shape,
        b_shape,
        w_init=w_init,
        b_init=b_init,
        w_regularizer=w_regularizer,
        trainable=trainable,
        name='main_gate')
    W_t, b_t = helper.weight_bias(
        w_shape,
        b_shape,
        w_init=w_init,
        b_init=b_init,
        w_regularizer=w_regularizer,
        trainable=trainable,
        name='transform_gate')
    H = activation(tf.matmul(x, W) + b, name='activation')
    T = tf.sigmoid(tf.matmul(x, W_t) + b_t, name='transform_gate')
    try:
      C = tf.sub(1.0, T, name="carry_gate")
    except Exception:
      C = tf.subtract(1.0, T, name="carry_gate")
    try:
      output = tf.add(tf.mul(H, T), tf.mul(x, C), name='output')
    except Exception:
      output = tf.add(tf.multiply(H, T), tf.multiply(x, C), name='output')

    if activation:
      output = activation(output, reuse=reuse, trainable=trainable)

    return _collect_named_outputs(outputs_collections, name, output)


def down_shifted_conv2d(x,
                        n_output_channels,
                        is_training,
                        reuse,
                        trainable=True,
                        filter_size=(2, 3),
                        stride=(1, 1),
                        **kwargs):
  """Down shifted convolution for PIXEL CNN."""
  x = tf.pad(
      x,
      [[0, 0], [filter_size[0] - 1, 0],
       [int((filter_size[1] - 1) / 2), int((filter_size[1] - 1) / 2)], [0, 0]])
  return conv2d(
      x,
      n_output_channels,
      is_training,
      reuse,
      trainable=trainable,
      filter_size=filter_size,
      padding='VALID',
      stride=stride,
      **kwargs)


def down_shifted_upsample2d(x,
                            output_shape,
                            is_training,
                            reuse,
                            trainable=True,
                            filter_size=(2, 3),
                            stride=(1, 1),
                            **kwargs):
  """Down shifted deconvolution for PIXEL CNN."""
  x = upsample2d(
      x,
      output_shape,
      is_training,
      reuse,
      trainable=trainable,
      filter_size=filter_size,
      pad='VALID',
      stride=stride,
      **kwargs)
  xs = list(map(int, x.get_shape()))
  return x[:, :(xs[1] - filter_size[0] + 1),
           int((filter_size[1] - 1) / 2):(xs[2] - int((filter_size[1] - 1) / 2)), :]


def down_right_shifted_conv2d(x,
                              n_output_channels,
                              is_training,
                              reuse,
                              filter_size=(2, 3),
                              stride=(1, 1),
                              **kwargs):
  """Down right shifted convolution for PIXEL CNN."""
  x = tf.pad(x, [[0, 0], [filter_size[0] - 1, 0], [filter_size[1] - 1, 0], [0, 0]])
  return conv2d(
      x,
      n_output_channels,
      is_training,
      reuse,
      trainable=trainable,
      filter_size=filter_size,
      padding='VALID',
      stride=stride,
      **kwargs)


def down_right_upsample2d(x,
                          output_shape,
                          is_training,
                          reuse,
                          trainable=True,
                          filter_size=(2, 2),
                          stride=(1, 1),
                          **kwargs):
  """Down right shifted deconvolution for PIXEL CNN."""
  x = upsample2d(
      x,
      output_shape,
      is_training,
      reuse,
      trainable=trainable,
      filter_size=filter_size,
      padding='VALID',
      stride=stride,
      **kwargs)
  xs = list(map(int, x.get_shape()))
  return x[:, :(xs[1] - filter_size[0] + 1):, :(xs[2] - filter_size[1] + 1), :]


def gated_resnet(x,
                 is_training,
                 reuse,
                 a=None,
                 h=None,
                 activation=tf.nn.relu,
                 conv=conv2d,
                 init=False,
                 counters={},
                 ema=None,
                 dropout_p=0.,
                 name='gated_resnet',
                 outputs_collections=None,
                 **kwargs):
  """Gated Resnet block."""
  with tf.name_scope(name):
    xs = list(map(int, x.get_shape()))
    num_filters = xs[-1]

    c1 = conv(activation(x), num_filters, is_training, reuse, **kwargs)
    if a is not None:
      c1 += conv(activation(a), num_filters, is_training, reuse, filter_size=(1, 1), **kwargs)
    c1 = activation(c1)
    if dropout_p > 0:
      c1 = dropout(c1, is_training, drop_p=dropout_p)
    c2 = conv(c1, num_filters * 2, is_training, reuse, init_scale=0.1, **kwargs)

    if h is not None:
      with tf.variable_scope(get_name('conditional_weights', counters)):
        hw = helper.get_var_maybe_avg(
            'hw',
            ema,
            shape=[int_shape(h)[-1], 2 * num_filters],
            dtype=tf.float32,
            initializer=tf.random_normal_initializer(0, 0.05),
            trainable=True)
      if init:
        hw = hw.initialized_value()
      c2 += tf.reshape(tf.matmul(h, hw), [xs[0], 1, 1, 2 * num_filters])

    a, b = tf.split(3, 2, c2)
    c3 = a * tf.nn.sigmoid(b)
    output = x + c3
    return _collect_named_outputs(outputs_collections, name, output)


def max_pool(x,
             filter_size=(3, 3),
             stride=(2, 2),
             padding='SAME',
             name='pool',
             outputs_collections=None,
             **unused):
  """Max pooling layer.

  Args:
      x: A 4-D 'Tensor` of shape `[batch_size, height, width, channels]`
      filter_size: A int or list/tuple of length 2: [kernel_height, kernel_width] of the
          pooling kernel over which the op is computed. Can be an int if both
          values are the same.
      stride: A int or list/tuple of length 2: [stride_height, stride_width].
      padding: The padding method, either 'VALID' or 'SAME'.
      outputs_collections: The collections to which the outputs are added.
      name: Optional scope/name for name_scope.

  Returns:
      A `Tensor` representing the results of the pooling operation.
      e.g.: 4-D `Tensor` [batch, new_height, new_width, channels].

  Raises:
      ValueError: If `input` is not 4-D array
  """
  _check_unused(unused, name)
  input_shape = helper.get_input_shape(x)
  assert len(input_shape) == 4, "Input Tensor shape must be 4-D"
  with tf.name_scope(name):
    output = tf.nn.max_pool(
        value=x,
        ksize=helper.kernel_2d(filter_size),
        strides=helper.stride_2d(stride),
        padding=helper.kernel_padding(padding),
    )
    return _collect_named_outputs(outputs_collections, name, output)


def max_pool_3d(x,
                filter_size=(3, 3, 3),
                stride=(2, 2, 2),
                padding='SAME',
                name='pool',
                outputs_collections=None,
                **unused):
  """Max pooling layer.

  Args:
      x: A 5-D 'Tensor` of shape `[batch_size, depth, height, width, channels]`
      filter_size: A int or list/tuple of length 3: [kernel_depth, kernel_height, kernel_width] of
          the pooling kernel over which the op is computed. Can be an int if both
          values are the same.
      stride: A int or list/tuple of length 3: [stride_depth, stride_height, stride_width].
      padding: The padding method, either 'VALID' or 'SAME'.
      outputs_collections: The collections to which the outputs are added.
      name: Optional scope/name for name_scope.

  Returns:
      A `Tensor` representing the results of the pooling operation.
      e.g.: 5-D `Tensor` [batch, new_depth, new_height, new_width, channels].

  Raises:
      ValueError: If `input` is not 5-D array
  """
  _check_unused(unused, name)
  input_shape = helper.get_input_shape(x)
  assert len(input_shape) == 5, "Input Tensor shape must be 4-D"
  with tf.name_scope(name):
    output = tf.nn.max_pool3d(
        x,
        helper.kernel_3d(filter_size),
        helper.stride_3d(stride),
        helper.kernel_padding(padding),
    )
    return _collect_named_outputs(outputs_collections, name, output)


def fractional_pool(x,
                    pooling_ratio=[1.0, 1.44, 1.73, 1.0],
                    pseudo_random=None,
                    determinastic=None,
                    overlapping=None,
                    name='fractional_pool',
                    seed=None,
                    seed2=None,
                    pool_type='avg',
                    outputs_collections=None,
                    **unused):
  """Fractional pooling layer.

  Args:
      x: A 4-D `Tensor` of shape `[batch_size, height, width, channels]`
      pooling_ratio: A list of floats that has length >= 4. Pooling ratio for each
          dimension of value, currently only supports row and col dimension and should
          be >= 1.0. For example, a valid pooling ratio looks like [1.0, 1.44, 1.73, 1.0].
          The first and last elements must be 1.0 because we don't allow pooling on batch and
          channels dimensions. 1.44 and 1.73 are pooling ratio on height and width
          dimensions respectively.
      pseudo_random: An optional bool. Defaults to False. When set to True, generates
          the pooling sequence in a pseudorandom fashion, otherwise, in a random fashion.
          Check paper Benjamin Graham, Fractional Max-Pooling for difference between
          pseudorandom and random.
      overlapping: An optional bool. Defaults to False. When set to True, it means when pooling,
          the values at the boundary of adjacent pooling cells are used by both cells.
          For example: index 0 1 2 3 4
          value 20 5 16 3 7; If the pooling sequence is [0, 2, 4], then 16, at index 2 will be used
          twice. The result would be [41/3, 26/3] for fractional avg pooling.
      deterministic: An optional bool. Defaults to False. When set to True, a fixed pooling
          region will be used when iterating over a FractionalAvgPool node in the computation
          graph. Mainly used in unit test to make FractionalAvgPool deterministic.
      seed: An optional int. Defaults to 0. If either seed or seed2 are set to be non-zero,
          the random number generator is seeded by the given seed. Otherwise,
          it is seeded by a random seed.
      seed2: An optional int. Defaults to 0. An second seed to avoid seed collision.
      outputs_collections: The collections to which the outputs are added.
      pool_type: avg or max pool
      name: Optional scope/name for name_scope.

  Returns:
      A 4-D `Tensor` representing the results of the pooling operation.
      e.g.: 4-D `Tensor` [batch, new_height, new_width, channels].

  Raises:
      ValueError: If 'kernel_size' is not a 2-D list
  """
  _check_unused(unused, name)
  input_shape = helper.get_input_shape(x)
  assert len(input_shape) == 4, "Input Tensor shape must be 4-D"
  with tf.name_scope(name):
    if pool_type == 'avg':
      outputs = tf.nn.fractional_avg_pool(
          x,
          pooling_ratio,
          pseudo_random=pseudo_random,
          overlapping=overlapping,
          deterministic=determinastic,
          seed=seed,
          seed2=seed2,
          name=name)
    else:
      outputs = tf.nn.fractional_max_pool(
          x,
          pooling_ratio,
          pseudo_random=pseudo_random,
          overlapping=overlapping,
          deterministic=determinastic,
          seed=seed,
          seed2=seed2,
          name=name)
    output = outputs[0]
    return _collect_named_outputs(outputs_collections, name, output)


def rms_pool_2d(x,
                filter_size=(3, 3),
                stride=(2, 2),
                padding='SAME',
                name='pool',
                epsilon=0.000000000001,
                outputs_collections=None,
                **unused):
  """RMS pooling layer.

  Args:
      x: A 4-D `Tensor` of shape `[batch_size, height, width, channels]`
      filter_size: A int or list/tuple of length 2: [kernel_height, kernel_width] of the
          pooling kernel over which the op is computed. Can be an int if both
          values are the same.
      stride: A int or list/tuple of length 2: [stride_height, stride_width].
      padding: The padding method, either 'VALID' or 'SAME'.
      outputs_collections: The collections to which the outputs are added.
      name: Optional scope/name for name_scope.
      epsilon: prevents divide by zero

  Returns:
      A 4-D `Tensor` representing the results of the pooling operation.
      e.g.: 4-D `Tensor` [batch, new_height, new_width, channels].

  Raises:
      ValueError: If 'input` not 4-D array
  """
  _check_unused(unused, name)
  input_shape = helper.get_input_shape(x)
  assert len(input_shape) == 4, "Input Tensor shape must be 4-D"
  with tf.name_scope(name):
    output = tf.nn.avg_pool(
        value=tf.square(x),
        ksize=helper.kernel_2d(filter_size),
        strides=helper.stride_2d(stride),
        padding=helper.kernel_padding(padding),
    )
    output = tf.sqrt(output + epsilon)
    return _collect_named_outputs(outputs_collections, name, output)


def rms_pool_3d(x,
                filter_size=(3, 3, 3),
                stride=(2, 2, 2),
                padding='SAME',
                name='pool',
                epsilon=0.000000000001,
                outputs_collections=None,
                **unused):
  """RMS pooling layer.

  Args:
      x: A 5-D `Tensor` of shape `[batch_size, depth, height, width, channels]`
      filter_size: A int or list/tuple of length 3: [kernel_depth, kernel_height, kernel_width] of
          the pooling kernel over which the op is computed. Can be an int if both
          values are the same.
      stride: A int or list/tuple of length 3: [stride_depth, stride_height, stride_width].
      padding: The padding method, either 'VALID' or 'SAME'.
      outputs_collections: The collections to which the outputs are added.
      name: Optional scope/name for name_scope.
      epsilon: prevents divide by zero

  Returns:
      A 5-D `Tensor` representing the results of the pooling operation.
      e.g.: 5-D `Tensor` [batch, new_height, new_width, channels].

  Raises:
      ValueError: If 'input' is not a 5-D array
  """
  _check_unused(unused, name)
  input_shape = helper.get_input_shape(x)
  assert len(input_shape) == 4, "Input Tensor shape must be 4-D"
  with tf.name_scope(name):
    output = tf.nn.avg_pool3d(
        tf.square(x),
        ksize=helper.kernel_3d(filter_size),
        strides=helper.stride_3d(stride),
        padding=helper.kernel_padding(padding),
    )
    output = tf.sqrt(output + epsilon)
    return _collect_named_outputs(outputs_collections, name, output)


def avg_pool_3d(x,
                filter_size=(3, 3, 3),
                stride=(2, 2, 2),
                padding='SAME',
                name=None,
                outputs_collections=None,
                **unused):
  """Avg pooling layer.

  Args:
      x: A 4-D `Tensor` of shape `[batch_size, depth, height, width, channels]`
      filter_size: A int or list/tuple of length 3: [kernel_depth, kernel_height, kernel_width] of
          the pooling kernel over which the op is computed. Can be an int if both
          values are the same.
      stride: A int or list/tuple of length 3: [stride_depth, stride_height, stride_width].
      padding: The padding method, either 'VALID' or 'SAME'.
      outputs_collections: The collections to which the outputs are added.
      name: Optional scope/name for name_scope.

  Returns:
      A 5-D `Tensor` representing the results of the pooling operation.
      e.g.: 5-D `Tensor` [batch, new_depth, new_height, new_width, channels].

  Raises:
      ValueError: If 'input' is not a 5-D array
  """
  _check_unused(unused, name)
  input_shape = helper.get_input_shape(x)
  assert len(input_shape) == 5, "Input Tensor shape must be 5-D"
  with tf.name_scope(name or "pool"):
    output = tf.nn.avg_pool3d(
        x,
        ksize=helper.kernel_3d(filter_size),
        strides=helper.stride_3d(stride),
        padding=helper.kernel_padding(padding),
        name="avg_pool")
    return _collect_named_outputs(outputs_collections, name, output)


def avg_pool_2d(x,
                filter_size=(3, 3),
                stride=(2, 2),
                padding='SAME',
                name=None,
                outputs_collections=None,
                **unused):
  """Avg pooling layer.

  Args:
      x: A 4-D `Tensor` of shape `[batch_size, height, width, channels]`
      filter_size: A int or list/tuple of length 2: [kernel_height, kernel_width] of the
          pooling kernel over which the op is computed. Can be an int if both
          values are the same.
      stride: A int or list/tuple of length 2: [stride_height, stride_width].
      padding: The padding method, either 'VALID' or 'SAME'.
      outputs_collections: The collections to which the outputs are added.
      name: Optional scope/name for name_scope.

  Returns:
      A 4-D `Tensor` representing the results of the pooling operation.
      e.g.: 4-D `Tensor` [batch, new_height, new_width, channels].

  Raises:
      ValueError: If 'kernel_size' is not a 2-D list
  """
  _check_unused(unused, name)
  input_shape = helper.get_input_shape(x)
  assert len(input_shape) == 4, "Input Tensor shape must be 4-D"
  with tf.name_scope(name or "pool"):
    output = tf.nn.avg_pool(
        value=x,
        ksize=helper.kernel_2d(filter_size),
        strides=helper.stride_2d(stride),
        padding=helper.kernel_padding(padding),
        name="avg_pool")
    return _collect_named_outputs(outputs_collections, name, output)


def global_avg_pool(x, name="global_avg_pool", outputs_collections=None, **unused):
  """Gloabl pooling layer.

  Args:
      x: A 4-D `Tensor` of shape `[batch_size, height, width, channels]`
      outputs_collections: The collections to which the outputs are added.
      name: Optional scope/name for name_scope.

  Returns:
      A 4-D `Tensor` representing the results of the pooling operation.
      e.g.: 4-D `Tensor` [batch, 1, 1, channels].

  Raises:
      ValueError: If 'kernel_size' is not a 2-D list
  """
  _check_unused(unused, name)
  input_shape = helper.get_input_shape(x)
  assert len(input_shape) == 4, "Input Tensor shape must be 4-D"
  with tf.name_scope(name):
    output = tf.reduce_mean(x, [1, 2])
    return _collect_named_outputs(outputs_collections, name, output)


def global_max_pool(x, name="global_max_pool", outputs_collections=None, **unused):
  """Gloabl max pooling layer.

  Args:
      x: A 4-D `Tensor` of shape `[batch_size, height, width, channels]`
      outputs_collections: The collections to which the outputs are added.
      name: Optional scope/name for name_scope.

  Returns:
      A 4-D `Tensor` representing the results of the pooling operation.
      e.g.: 4-D `Tensor` [batch, 1, 1, channels].

  Raises:
      ValueError: If 'kernel_size' is not a 2-D list
  """
  _check_unused(unused, name)
  input_shape = helper.get_input_shape(x)
  assert len(input_shape) == 4, "Input Tensor shape must be 4-D"
  with tf.name_scope(name):
    output = tf.reduce_max(x, [1, 2])
    return _collect_named_outputs(outputs_collections, name, output)


def feature_max_pool_1d(x, stride=2, name='feature_max_pool_1d', outputs_collections=None, **unused):
  """Feature max pooling layer.

  Args:
      x: A 2-D tensor of shape `[batch_size, channels]`
      stride: A int.
      outputs_collections: The collections to which the outputs are added.
      name: Optional scope/name for name_scope.

  Returns:
      A 2-D `Tensor` representing the results of the pooling operation.
      e.g.: 2-D `Tensor` [batch_size, new_channels]

  Raises:
      ValueError: If 'kernel_size' is not a 2-D list
  """
  _check_unused(unused, name)
  input_shape = helper.get_input_shape(x)
  assert len(input_shape) == 2, "Input Tensor shape must be 2-D"
  x = tf.reshape(x, (-1, input_shape[1] // stride, stride))
  with tf.name_scope(name):
    output = tf.reduce_max(
        input_tensor=x,
        reduction_indices=[2],
    )
    return _collect_named_outputs(outputs_collections, name, output)


def feature_max_pool_2d(x, stride=2, name='feature_max_pool_2d', outputs_collections=None, **unused):
  """Feature max pooling layer.

  Args:
      x: A 4-D tensor of shape `[batch_size, height, width, channels]`
      stride: A int.
      outputs_collections: The collections to which the outputs are added.
      name: Optional scope/name for name_scope.

  Returns:
      A 4-D `Tensor` representing the results of the pooling operation.
      e.g.: 4-D `Tensor` [batch_size, height, width, new_channels]

  Raises:
      ValueError: If 'kernel_size' is None
  """
  _check_unused(unused, name)
  input_shape = helper.get_input_shape(x)
  assert len(input_shape) == 4, "Input Tensor shape must be 4-D"
  x = tf.reshape(x, (-1, input_shape[1], input_shape[2], input_shape[3] // stride, stride))
  with tf.name_scope(name):
    output = tf.reduce_max(
        input_tensor=x,
        reduction_indices=[4],
    )
    return _collect_named_outputs(outputs_collections, name, output)


def local_response_normalization(x,
                                 depth_radius=5,
                                 bias=1,
                                 alpha=1,
                                 beta=0.5,
                                 name='local_response_normalization',
                                 outputs_collections=None,
                                 **unused):
  _check_unused(unused, name)
  input_shape = helper.get_input_shape(x)
  assert len(input_shape) == 4, "Input Tensor shape must be 4-D"
  with tf.name_scope(name):
    output = tf.nn.local_response_normalization(
        input=x, depth_radius=depth_radius, bias=bias, alpha=alpha, beta=beta)
    return _collect_named_outputs(outputs_collections, name, output)


def batch_norm_tf(x, name='bn', scale=False, fused=False, is_training=False, decay=0.999, **kwargs):
  """Adds a Batch Normalization layer from http://arxiv.org/abs/1502.03167.
        "Batch Normalization: Accelerating Deep Network Training by Reducing
        Internal Covariate Shift", Sergey Ioffe, Christian Szegedy
        Can be used as a normalizer function for conv2d and fully_connected.
        Note: When is_training is True the moving_mean and moving_variance need to be
        updated, by default the update_ops are placed in `tf.GraphKeys.UPDATE_OPS` so
        they need to be added as a dependency to the `train_op`, example:
        `update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)`
        `if update_ops:`
            `updates = tf.group(*update_ops)`
        `total_loss = control_flow_ops.with_dependencies([updates], total_loss)`
        One can set updates_collections=None to force the updates in place, but that
        can have speed penalty, specially in distributed settings.
    Args:
        x: a `Tensor` with 2 or more dimensions, where the first dimension has
            `batch_size`. The normalization is over all but the last dimension if
            `data_format` is `NHWC` and the second dimension if `data_format` is
            `NCHW`.
        decay: decay for the moving average. Reasonable values for `decay` are close
            to 1.0, typically in the multiple-nines range: 0.999, 0.99, 0.9, etc.
            Lower `decay` value (recommend trying `decay`=0.9) if model experiences
            reasonably good training performance but poor validation and/or test
            performance. Try zero_debias_moving_mean=True for improved stability.
        center: If True, subtract `beta`. If False, `beta` is ignored.
        scale: If True, multiply by `gamma`. If False, `gamma` is
            not used. When the next layer is linear (also e.g. `nn.relu`), this can be
            disabled since the scaling can be done by the next layer.
        epsilon: small float added to variance to avoid dividing by zero.
        activation_fn: activation function, default set to None to skip it and
            maintain a linear activation.
        param_initializers: optional initializers for beta, gamma, moving mean and
            moving variance.
        updates_collections: collections to collect the update ops for computation.
            The updates_ops need to be executed with the train_op.
            If None, a control dependency would be added to make sure the updates are
            computed in place.
        is_training: whether or not the layer is in training mode. In training mode
            it would accumulate the statistics of the moments into `moving_mean` and
            `moving_variance` using an exponential moving average with the given
            `decay`. When it is not in training mode then it would use the values of
            the `moving_mean` and the `moving_variance`.
        reuse: whether or not the layer and its variables should be reused. To be
            able to reuse the layer scope must be given.
            outputs_collections: collections to add the outputs.
        trainable: If `True` also add variables to the graph collection
            `GraphKeys.TRAINABLE_VARIABLES` (see `tf.Variable`).
        batch_weights: An optional tensor of shape `[batch_size]`,
            containing a frequency weight for each batch item. If present,
            then the batch normalization uses weighted mean and
            variance. (This can be used to correct for bias in training
            example selection.)
        fused:  Use nn.fused_batch_norm if True, nn.batch_normalization otherwise.
        name: Optional scope/name for `variable_scope`.

    Returns:
        A `Tensor` representing the output of the operation.

    Raises:
        ValueError: if `batch_weights` is not None and `fused` is True.
        ValueError: if the rank of `inputs` is undefined.
        ValueError: if rank or channels dimension of `inputs` is undefined.
    """
  return tf.layers.batch_normalization(
      x, name=name, scale=scale, training=is_training, fused=fused, momentum=decay, **kwargs)


def batch_norm_lasagne(x,
                       is_training,
                       reuse,
                       trainable=True,
                       decay=0.9,
                       epsilon=1e-4,
                       name='bn',
                       updates_collections=tf.GraphKeys.UPDATE_OPS,
                       outputs_collections=None):
  """Adds a Batch Normalization layer from http://arxiv.org/abs/1502.03167.
        Instead of storing and updating moving variance, this layer store and
        update moving inverse standard deviation
        "Batch Normalization: Accelerating Deep Network Training by Reducin Internal Covariate Shift"
        Sergey Ioffe, Christian Szegedy
        Can be used as a normalizer function for conv2d and fully_connected.
        Note: When is_training is True the moving_mean and moving_variance need to be
        updated, by default the update_ops are placed in `tf.GraphKeys.UPDATE_OPS` so
        they need to be added as a dependency to the `train_op`, example:
        `update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)`
        `if update_ops:`
            `updates = tf.group(*update_ops)`
        `total_loss = control_flow_ops.with_dependencies([updates], total_loss)`
        One can set updates_collections=None to force the updates in place, but that
        can have speed penalty, specially in distributed settings.

    Args:
        x: a tensor with 2 or more dimensions, where the first dimension has
            `batch_size`. The normalization is over all but the last dimension if
            `data_format` is `NHWC` and the second dimension if `data_format` is
            `NCHW`.
        decay: decay for the moving average. Reasonable values for `decay` are close
            to 1.0, typically in the multiple-nines range: 0.999, 0.99, 0.9, etc.
            Lower `decay` value (recommend trying `decay`=0.9) if model experiences
            reasonably good training performance but poor validation and/or test
            performance. Try zero_debias_moving_mean=True for improved stability.
        epsilon: small float added to variance to avoid dividing by zero.
        updates_collections: collections to collect the update ops for computation.
            The updates_ops need to be executed with the train_op.
            If None, a control dependency would be added to make sure the updates are
            computed in place.
        is_training: whether or not the layer is in training mode. In training mode
            it would accumulate the statistics of the moments into `moving_mean` and
            `moving_variance` using an exponential moving average with the given
            `decay`. When it is not in training mode then it would use the values of
            the `moving_mean` and the `moving_variance`.
        reuse: whether or not the layer and its variables should be reused. To be
            able to reuse the layer scope must be given.
        outputs_collections: collections to add the outputs.
        trainable: If `True` also add variables to the graph collection
            `GraphKeys.TRAINABLE_VARIABLES` (see `tf.Variable`).
        name: Optional scope/name for `variable_scope`.

    Returns:
        A `Tensor` representing the output of the operation.

    Raises:
        ValueError: if the rank of `x` is undefined.
        ValueError: if rank or channels dimension of `inputs` is undefined.
    """
  with tf.variable_scope(name, reuse=reuse):
    beta = tf.get_variable(
        name='beta', initializer=tf.constant(0.0, shape=[x.get_shape()[-1]]), trainable=trainable)
    gamma = tf.get_variable(
        name='gamma', initializer=tf.constant(1.0, shape=[x.get_shape()[-1]]), trainable=trainable)

    moving_mean = tf.get_variable(
        name='moving_mean',
        shape=[x.get_shape()[-1]],
        initializer=tf.zeros_initializer(dtype=tf.float32),
        trainable=False)

    moving_inv_std = tf.get_variable(
        name='moving_inv_std',
        shape=[x.get_shape()[-1]],
        initializer=tf.ones_initializer(dtype=tf.float32),
        trainable=False)

    def mean_inv_std_with_update():
      mean, variance = tf.nn.moments(x, [0, 1, 2], shift=moving_mean, name='bn-moments')
      inv_std = math_ops.rsqrt(variance + epsilon)
      update_moving_mean = moving_averages.assign_moving_average(
          moving_mean, mean, decay, zero_debias=False)
      update_moving_inv_std = moving_averages.assign_moving_average(
          moving_inv_std, inv_std, decay, zero_debias=False)
      with tf.control_dependencies([update_moving_mean, update_moving_inv_std]):
        m, v = tf.identity(mean), tf.identity(inv_std)
        return m, v

    def mean_inv_std_with_pending_update():
      mean, variance = tf.nn.moments(x, [0, 1, 2], shift=moving_mean, name='bn-moments')
      inv_std = math_ops.rsqrt(variance + epsilon)
      update_moving_mean = moving_averages.assign_moving_average(
          moving_mean, mean, decay, zero_debias=False)
      update_moving_inv_std = moving_averages.assign_moving_average(
          moving_inv_std, inv_std, decay, zero_debias=False)
      tf.add_to_collection(updates_collections, update_moving_mean)
      tf.add_to_collection(updates_collections, update_moving_inv_std)
      return mean, inv_std

    mean_inv_std_with_relevant_update = mean_inv_std_with_pending_update if updates_collections \
        is not None else mean_inv_std_with_update

    (mean, inv_std) = mean_inv_std_with_relevant_update() if is_training else (moving_mean,
                                                                               moving_inv_std)

    def _batch_normalization(x, mean, inv, offset, scale):
      with tf.name_scope(name, "batchnorm", [x, mean, inv, scale, offset]):
        if scale is not None:
          inv *= scale
        return x * inv + (offset - mean * inv if offset is not None else -mean * inv)

    output = _batch_normalization(x, mean, inv_std, beta, gamma)
    return _collect_named_outputs(outputs_collections, name, output)


def layer_norm(x,
               reuse,
               filters=None,
               trainable=True,
               epsilon=1e-6,
               name='layer_norm',
               allow_defun=False,
               outputs_collections=None):
  """Layer normalize the tensor x, averaging over the last dimension
    Args:
        x: a `Tensor` with type `float`, `double`, `int32`, `int64`, `uint8`, int16`, or `int8`.
        reuse: whether or not the layer and its variables should be reused. To be
            able to reuse the layer scope must be given.
        trainable: a bool, training or fixed value
        name: a optional scope/name of the layer
        outputs_collections: The collections to which the outputs are added.

    Returns:
        A `Tensor` representing the results of the layer norm operation.
    """
  if filters is None:
    filters = x.get_shape()[-1]
  with tf.variable_scope(name, reuse=reuse):
    scale = tf.get_variable(
        "layer_norm_scale", [filters], initializer=tf.ones_initializer(), trainable=trainable)
    bias = tf.get_variable(
        "layer_norm_bias", [filters], initializer=tf.zeros_initializer(), trainable=trainable)
    if allow_defun:
      output = layer_norm_compute(x, tf.constant(epsilon), scale, bias)
      output.set_shape(x.get_shape())
    else:
      output = _layer_norm_compute_python(x, epsilon, scale, bias)

  return _collect_named_outputs(outputs_collections, name, output)


def _layer_norm_compute_python(x, epsilon, scale, bias):
  """Layer norm raw computation."""
  mean = tf.reduce_mean(x, axis=[-1], keep_dims=True)
  variance = tf.reduce_mean(tf.square(x - mean), axis=[-1], keep_dims=True)
  norm_x = (x - mean) * tf.rsqrt(variance + epsilon)

  return norm_x * scale + bias


@function.Defun(compiled=True)
def layer_norm_compute_grad(x, epsilon, scale, bias, dy):
  y = _layer_norm_compute_python(x, epsilon, scale, bias)
  dx = tf.gradients(ys=[y], xs=[x, epsilon, scale, bias], grad_ys=[dy])
  return dx


@function.Defun(compiled=True, separate_compiled_gradients=True, grad_func=layer_norm_compute_grad)
def layer_norm_compute(x, epsilon, scale, bias):
  return _layer_norm_compute_python(x, epsilon, scale, bias)


def prelu(x, reuse, alpha_init=0.2, trainable=True, name='prelu', outputs_collections=None):
  """Prametric rectifier linear layer.

  Args:
      x: a `Tensor` with type `float`, `double`, `int32`, `int64`, `uint8`, int16`, or `int8`.
      reuse: whether or not the layer and its variables should be reused. To be
          able to reuse the layer scope must be given.
      alpha_init: initalization value for alpha
      trainable: a bool, training or fixed value
      name: a optional scope/name of the layer
      outputs_collections: The collections to which the outputs are added.

  Returns:
      A `Tensor` representing the results of the prelu activation operation.
  """
  with tf.variable_scope(name, reuse=reuse):
    alphas = tf.get_variable(
        name='alpha',
        initializer=tf.constant(alpha_init, shape=[x.get_shape()[-1]]),
        trainable=trainable)
    try:
      output = tf.nn.relu(x) + tf.mul(alphas, (x - tf.abs(x))) * 0.5
    except Exception:
      output = tf.nn.relu(x) + tf.multiply(alphas, (x - tf.abs(x))) * 0.5
    return _collect_named_outputs(outputs_collections, name, output)


def relu(x, name='relu', outputs_collections=None, **unused):
  """Computes relu.

  Args:
      x: a `Tensor` with type `float`, `double`, `int32`, `int64`, `uint8`, int16`, or `int8`.
      name: a optional scope/name of the layer
      outputs_collections: The collections to which the outputs are added.

  Returns:
      A `Tensor` representing the results of the activation operation.
  """
  _check_unused(unused, name)
  with tf.name_scope(name):
    output = tf.nn.relu(x)
    return _collect_named_outputs(outputs_collections, name, output)


def relu6(x, name='relu6', outputs_collections=None, **unused):
  """Rectifier linear relu6 layer.

  Args:
      x: a `Tensor` with type `float`, `double`, `int32`, `int64`, `uint8`, int16`, or `int8`.
      name: a optional scope/name of the layer
      outputs_collections: The collections to which the outputs are added.

  Returns:
      A `Tensor` representing the results of the relu6 activation operation.
  """
  _check_unused(unused, name)
  with tf.name_scope(name):
    output = tf.nn.relu6(x)
    return _collect_named_outputs(outputs_collections, name, output)


def softplus(x, name='softplus', outputs_collections=None, **unused):
  """Softplus layer Computes softplus: log(exp(x) + 1).

  Args:
      x: a `Tensor` with type `float`, `double`, `int32`, `int64`, `uint8`, int16`, or `int8`.
      name: a optional scope/name of the layer
      outputs_collections: The collections to which the outputs are added.

  Returns:
      A `Tensor` representing the results of the activation operation.
  """
  _check_unused(unused, name)
  with tf.name_scope(name):
    output = tf.nn.softplus(x)
    return _collect_named_outputs(outputs_collections, name, output)


def softsign(x, name='softsign', outputs_collections=None, **unused):
  """Softsign layer Computes softsign: x / (abs(x) + 1).

  Args:
      x: a `Tensor` with type `float`, `double`, `int32`, `int64`, `uint8`, int16`, or `int8`.
      name: a optional scope/name of the layer
      outputs_collections: The collections to which the outputs are added.

  Returns:
      A `Tensor` representing the results of the activation operation.
  """
  _check_unused(unused, name)
  with tf.name_scope(name):
    output = tf.nn.softsign(x)
    return _collect_named_outputs(outputs_collections, name, output)


def crelu(x, name='crelu', outputs_collections=None, **unused):
  """Computes Concatenated ReLU.
        Concatenates a ReLU which selects only the positive part of the activation with
        a ReLU which selects only the negative part of the activation. Note that
        at as a result this non-linearity doubles the depth of the activations.
        Source: https://arxiv.org/abs/1603.05201

    Args:
        x: a `Tensor` with type `float`, `double`, `int32`, `int64`, `uint8`, int16`, or `int8`.
        name: a optional scope/name of the layer
        outputs_collections: The collections to which the outputs are added.

    Returns:
        A `Tensor` representing the results of the activation operation.
    """
  _check_unused(unused, name)
  with tf.name_scope(name):
    output = tf.nn.crelu(x)
    return _collect_named_outputs(outputs_collections, name, output)


def elu(x, alpha=1.0, name='elu', outputs_collections=None, **unused):
  """Computes exponential linear: exp(features) - 1 if < 0, features otherwise.
        See "Fast and Accurate Deep Network Learning by Exponential Linear Units (ELUs)"

    Args:
        x: a `Tensor` with type `float`, `double`, `int32`, `int64`, `uint8`, int16`, or `int8`.
        name: a optional scope/name of the layer
        outputs_collections: The collections to which the outputs are added.

    Returns:
        A `Tensor` representing the results of the activation operation.
    """
  _check_unused(unused, name)
  with tf.name_scope(name):
    output = tf.nn.elu(x)
    if alpha != 1.0:
      output = tf.where(output > 0, output, alpha * output)
    return _collect_named_outputs(outputs_collections, name, output)


def concat_elu(x, name='concat_elu', outputs_collections=None, **unused):
  """like concatenated ReLU (http://arxiv.org/abs/1603.05201), but then with
  ELU.

  Args:
      x: a `Tensor` with type `float`, `double`, `int32`, `int64`, `uint8`, int16`, or `int8`.
      name: a optional scope/name of the layer
      outputs_collections: The collections to which the outputs are added.

  Returns:
      A `Tensor` representing the results of the activation operation.
  """
  with tf.name_scope(name):
    axis = len(x.get_shape()) - 1
    output = tf.nn.elu(tf.concat([x, -x], axis=axis))
    return _collect_named_outputs(outputs_collections, name, output)


def leaky_relu(x, alpha=0.01, name='leaky_relu', outputs_collections=None, **unused):
  """Computes leaky relu.

  Args:
      x: a `Tensor` with type `float`, `double`, `int32`, `int64`, `uint8`, int16`, or `int8`.
      aplha: the conatant fro scalling the activation
      name: a optional scope/name of the layer
      outputs_collections: The collections to which the outputs are added.

  Returns:
      A `Tensor` representing the results of the activation operation.
  """
  _check_unused(unused, name)
  with tf.name_scope(name):
    try:
      output = tf.nn.relu(x) + tf.mul(alpha, (x - tf.abs(x))) * 0.5
    except Exception:
      output = tf.nn.relu(x) + tf.multiply(alpha, (x - tf.abs(x))) * 0.5
    return _collect_named_outputs(outputs_collections, name, output)


def lrelu(x, leak=0.2, name="lrelu", outputs_collections=None, **unused):
  """Computes reaky relu lasagne style.

  Args:
      x: a `Tensor` with type `float`, `double`, `int32`, `int64`, `uint8`, int16`, or `int8`.
      leak: the conatant fro scalling the activation
      name: a optional scope/name of the layer
      outputs_collections: The collections to which the outputs are added.

  Returns:
      A `Tensor` representing the results of the activation operation.
  """
  with tf.variable_scope(name):
    f1 = 0.5 * (1 + leak)
    f2 = 0.5 * (1 - leak)
    output = f1 * x + f2 * abs(x)
    return _collect_named_outputs(outputs_collections, name, output)


def maxout(x, k=2, name='maxout', outputs_collections=None, **unused):
  """Computes maxout activation.

  Args:
      x: a `Tensor` with type `float`, `double`, `int32`, `int64`, `uint8`, int16`, or `int8`.
      k: output channel splitting factor
      name: a optional scope/name of the layer
      outputs_collections: The collections to which the outputs are added.

  Returns:
      A `Tensor` representing the results of the activation operation.
  """
  with tf.name_scope(name):
    shape = [int(e) for e in x.get_shape()]
    ax = len(shape)
    ch = shape[-1]
    assert ch % k == 0
    shape[-1] = ch / k
    shape.append(k)
    x = tf.reshape(x, shape)
    output = tf.reduce_max(x, ax)
    return _collect_named_outputs(outputs_collections, name, output)


def offset_maxout(x, k=2, name='maxout', outputs_collections=None, **unused):
  """Computes maxout activation.

  Args:
      x: a `Tensor` with type `float`, `double`, `int32`, `int64`, `uint8`, int16`, or `int8`.
      k: output channel splitting factor
      name: a optional scope/name of the layer
      outputs_collections: The collections to which the outputs are added.

  Returns:
      A `Tensor` representing the results of the activation operation.
  """
  with tf.name_scope(name):
    shape = [int(e) for e in x.get_shape()]
    ax = len(shape)
    ch = shape[-1]
    assert ch % k == 0
    shape[-1] = ch / k
    shape.append(k)
    x = tf.reshape(x, shape)
    ofs = rng.randn(1000, k).max(axis=1).mean()
    output = tf.reduce_max(x, ax) - ofs
    return _collect_named_outputs(outputs_collections, name, output)


def sigmoid(x, name='softmax', outputs_collections=None, **unused):
  """Computes sigmoid activation.

  Args:
      x: a `Tensor` with type `float`, `double`, `int32`, `int64`, `uint8`, int16`, or `int8`.
      name: a optional scope/name of the layer
      outputs_collections: The collections to which the outputs are added.

  Returns:
      A `Tensor` representing the results of the activation operation.
  """
  _check_unused(unused, name)
  with tf.name_scope(name):
    output = tf.sigmoid(x)
    return _collect_named_outputs(outputs_collections, name, output)


def softmax(x, name='softmax', outputs_collections=None, **unused):
  """Computes softmax activation.

  Args:
      x: a `Tensor` with type `float`, `double`, `int32`, `int64`, `uint8`, int16`, or `int8`.
      name: a optional scope/name of the layer
      outputs_collections: The collections to which the outputs are added.

  Returns:
      A `Tensor` representing the results of the activation operation.
  """
  _check_unused(unused, name)
  with tf.name_scope(name):
    output = tf.nn.softmax(x)
    return _collect_named_outputs(outputs_collections, name, output)


def spatial_softmax(features,
                    reuse,
                    temperature=None,
                    name='spatial_softmax',
                    trainable=True,
                    outputs_collections=None,
                    **unused):
  """Computes the spatial softmax of a convolutional feature map. First
  computes the softmax over the spatial extent of each channel of a
  convolutional feature map. Then computes the expected 2D position of the
  points of maximal activation for each channel, resulting in a set of feature
  keypoints [x1, y1, ... xN, yN] for all N channels. Read more here: "Learning
  visual feature spaces for robotic manipulation with deep spatial
  autoencoders." Finn et. al, http://arxiv.org/abs/1509.06113.

  Args:
      features: A `Tensor` of size [batch_size, W, H, num_channels]; the
          convolutional feature map.
      reuse: whether or not the layer and its variables should be reused. To be
          able to reuse the layer scope must be given.
      outputs_collections: The collections to which the outputs are added.
      temperature: Softmax temperature (optional). If None, a learnable
          temperature is created.
      name: A name for this operation (optional).
      trainable: If `True` also add variables to the graph collection
          `GraphKeys.TRAINABLE_VARIABLES` (see `tf.Variable`).

  Returns:
      feature_keypoints: A `Tensor` with size [batch_size, num_channels * 2];
          the expected 2D locations of each channel's feature keypoint (normalized
          to the range (-1,1)). The inner dimension is arranged as
         [x1, y1, ... xN, yN].

  Raises:
      ValueError: If unexpected data_format specified.
      ValueError: If num_channels dimension is unspecified.
  """
  shape = tf.shape(features)
  static_shape = features.shape
  height, width, num_channels = shape[1], shape[2], static_shape[3]
  if num_channels.value is None:
    raise ValueError('The num_channels dimension of the inputs to '
                     '`spatial_softmax` should be defined. Found `None`.')

  with tf.name_scope(name, 'spatial_softmax', [features]) as name:
    # Create tensors for x and y coordinate values, scaled to range [-1, 1].
    pos_x, pos_y = tf.meshgrid(
        tf.lin_space(-1., 1., num=height), tf.lin_space(-1., 1., num=width), indexing='ij')
    pos_x = tf.reshape(pos_x, [height * width])
    pos_y = tf.reshape(pos_y, [height * width])
    if temperature is None:
      with tf.variable_scope(name + '_temperature', reuse=reuse):
        temperature = tf.get_variable(
            'temperature',
            shape=(),
            dtype=tf.float32,
            initializer=tf.ones_initializer(),
            trainable=trainable)
    features = tf.reshape(tf.transpose(features, [0, 3, 1, 2]), [-1, height * width])

    softmax_attention = tf.nn.softmax(features / temperature)
    expected_x = tf.reduce_sum(pos_x * softmax_attention, [1], keep_dims=True)
    expected_y = tf.reduce_sum(pos_y * softmax_attention, [1], keep_dims=True)
    expected_xy = tf.concat([expected_x, expected_y], 1)
    feature_keypoints = tf.reshape(expected_xy, [-1, num_channels.value * 2])
    feature_keypoints.set_shape([None, num_channels.value * 2])

    return _collect_named_outputs(outputs_collections, name, feature_keypoints)


def selu(x, alpha=None, scale=None, name='selu', outputs_collections=None, **unused):
  """Computes selu.

  Args:
      x: a `Tensor` with type `float`, `double`, `int32`, `int64`, `uint8`, int16`, or `int8`.
      alpha: float, selu parameters calculated from fixed points
      scale: float, selu parameters calculated from fixed points
      name: a optional scope/name of the layer
      outputs_collections: The collections to which the outputs are added.

  Returns:
      A `Tensor` representing the results of the selu activation operation.
  """
  _check_unused(unused, name)
  with tf.name_scope(name):
    if None in (alpha, scale):
      # using parameters from 0 mean, unit variance points
      alpha = 1.6732632423543772848170429916717
      scale = 1.0507009873554804934193349852946
    output = scale * tf.where(x >= 0.0, x, alpha * tf.nn.elu(x))
    return _collect_named_outputs(outputs_collections, name, output)


def dropout_selu(x,
                 is_training,
                 drop_p=0.2,
                 alpha=-1.7580993408473766,
                 fixedPointMean=0.0,
                 fixedPointVar=1.0,
                 noise_shape=None,
                 seed=None,
                 name='dropout_selu',
                 outputs_collections=None,
                 **unused):
  """
    Dropout layer for self normalizing networks
    Args:
        x: a `Tensor`.
        is_training: a bool, training or validation
        drop_p: probability of droping unit
        fixedPointsMean: float, the mean used to calculate the selu parameters
        fixedPointsVar: float, the Variance used to calculate the selu parameters
        alpha: float, product of the two selu parameters
        name: a optional scope/name of the layer
        outputs_collections: The collections to which the outputs are added.

    Returns:
        A `Tensor` representing the results of the dropout operation.
    """
  _check_unused(unused, name)

  def dropout_selu_impl(x, drop_p, alpha, noise_shape, seed, name):
    keep_prob = 1.0 - drop_p
    if isinstance(keep_prob, numbers.Real) and not 0 < keep_prob <= 1:
      raise ValueError("keep_prob must be a scalar tensor or a float in the "
                       "range (0, 1], got %g" % keep_prob)
    keep_prob = tf.convert_to_tensor(keep_prob, dtype=x.dtype, name="keep_prob")
    keep_prob.get_shape().assert_has_rank(0)

    alpha = tf.convert_to_tensor(alpha, dtype=x.dtype, name="alpha")
    alpha.get_shape().assert_has_rank(0)

    if tf.contrib.util.constant_value(keep_prob) == 1:
      return x

    noise_shape = noise_shape if noise_shape is not None else x.get_shape().as_list()
    random_tensor = keep_prob
    random_tensor += tf.random_uniform(noise_shape, seed=seed, dtype=x.dtype)
    binary_tensor = tf.floor(random_tensor)
    ret = x * binary_tensor + alpha * (1 - binary_tensor)

    a = tf.sqrt(fixedPointVar / (keep_prob * (
        (1 - keep_prob) * math_ops.pow(alpha - fixedPointMean, 2) + fixedPointVar)))

    b = fixedPointMean - a * \
        (keep_prob * fixedPointMean + (1 - keep_prob) * alpha)
    ret = a * ret + b
    ret.set_shape(x.get_shape())
    return ret

  with tf.name_scope(name, [x]):
    if is_training:
      output = dropout_selu_impl(x, drop_p, alpha, noise_shape, seed, name)
    else:
      output = x
    return _collect_named_outputs(outputs_collections, name, output)


def sample_gumbel(shape, eps=1e-20):
  """Sample from Gumbel(0, 1)"""
  U = tf.random_uniform(shape, minval=0, maxval=1)
  return -tf.log(-tf.log(U + eps) + eps)


def gumbel_softmax_sample(logits, temperature):
  """Draw a sample from the Gumbel-Softmax distribution."""
  y = logits + sample_gumbel(tf.shape(logits))
  return tf.nn.softmax(y / temperature)


def gumbel_softmax(logits, temperature, hard=False):
  """Computes Gumbel Softmax Sample from the Gumbel-Softmax distribution and
  optionally discretize. http://blog.evjang.com/2016/11/tutorial-categorical-
  variational.html https://arxiv.org/abs/1611.01144.

  Args:
      logits: [batch_size, n_class] unnormalized log-probs
      temperature: non-negative scalar
      hard: if True, take argmax, but differentiate w.r.t. soft sample y

  Returns:
      [batch_size, n_class] sample from the Gumbel-Softmax distribution.
      If hard=True, then the returned sample will be one-hot, otherwise it will
      be a probabilitiy distribution that sums to 1 across classes
  """
  y = gumbel_softmax_sample(logits, temperature)
  if hard:
    k = tf.shape(logits)[-1]
    # y_hard = tf.cast(tf.one_hot(tf.argmax(y,1),k), y.dtype)
    y_hard = tf.cast(tf.equal(y, tf.reduce_max(y, 1, keep_dims=True)), y.dtype)
    y = tf.stop_gradient(y_hard - y) + y
  return y


def pixel_wise_softmax(inputs):
  """Computes pixel wise softmax activation.

  Args:
      x: a `Tensor` with type `float`, `double`, `int32`, `int64`, `uint8`, int16`, or `int8`.
      name: a optional scope/name of the layer
      outputs_collections: The collections to which the outputs are added.

  Returns:
      A `Tensor` representing the results of the activation operation.
  """
  exponential_map = tf.exp(inputs)
  sum_exp = tf.reduce_sum(exponential_map, 3, keep_dims=True)
  tensor_sum_exp = tf.tile(sum_exp, tf.stack([1, 1, 1, tf.shape(inputs)[3]]))
  return tf.div(exponential_map, tensor_sum_exp)


def dropout(x,
            is_training,
            drop_p=0.5,
            seed=None,
            name='dropout',
            outputs_collections=None,
            **unused):
  """
    Dropout layer
    Args:
        x: a `Tensor`.
        is_training: a bool, training or validation
        drop_p: probability of droping unit
        name: a optional scope/name of the layer
        outputs_collections: The collections to which the outputs are added.

    Returns:
        A `Tensor` representing the results of the dropout operation.
    """
  _check_unused(unused, name)
  with tf.name_scope(name):
    keep_p = 1. - drop_p
    if is_training:
      output = tf.nn.dropout(x, keep_p, seed=seed)
    else:
      output = x
    return _collect_named_outputs(outputs_collections, name, output)


def _flatten(x, name='flatten'):
  input_shape = helper.get_input_shape(x)
  assert len(input_shape) > 1, "Input Tensor shape must be > 1-D"
  with tf.name_scope(name):
    dims = int(np.prod(input_shape[1:]))
    flattened = tf.reshape(x, [-1, dims])
    return flattened


def flatten(x, name='flatten'):
  input_shape = helper.get_input_shape(x)
  assert len(input_shape) > 1, "Input Tensor shape must be > 1-D"
  with tf.name_scope(name):
    dims = int(np.prod(input_shape[1:]))
    flattened = tf.reshape(x, [-1, dims])
    return flattened


def repeat(x,
           repetitions,
           layer,
           num_outputs=None,
           name='Repeat',
           outputs_collections=None,
           *args,
           **kwargs):
  """Repeat op.

  Args:
      x: a `Tensor`.
      repetitions: a int, number of times to apply the same operation
      layer: the layer function with arguments to repeat
      name: a optional scope/name of the layer
      outputs_collections: The collections to which the outputs are added.

  Returns:
      A `Tensor` representing the results of the repetition operation.
  """
  with tf.variable_scope(name, 'Repeat'):
    inputs = tf.convert_to_tensor(x)
    if name is None:
      if hasattr(layer, '__name__'):
        name = layer.__name__
      elif hasattr(layer, 'func') and hasattr(layer.func, '__name__'):
        # In case layer is a functools.partial.
        name = layer.func.__name__
      else:
        name = 'Repeat'
    outputs = inputs
    for i in range(repetitions):
      new_name = name + '_' + str(i + 1)
      if num_outputs is None:
        outputs = layer(outputs, name=new_name, *args, **kwargs)
      else:
        outputs = layer(outputs, num_outputs, name=new_name, *args, **kwargs)
      tf.add_to_collection(outputs_collections, NamedOutputs(new_name, outputs))
    return outputs


def merge(tensors_list, mode, axis=1, name='merge', outputs_collections=None, **kwargs):
  """Merge op.

  Args:
      tensor_list: A list `Tensors` to merge
      mode: str, available modes are
          ['concat', 'elemwise_sum', 'elemwise_mul', 'sum',
              'mean', 'prod', 'max', 'min', 'and', 'or']
      name: a optional scope/name of the layer
      outputs_collections: The collections to which the outputs are added.

  Returns:
      A `Tensor` representing the results of the repetition operation.

  Raises:
      ValueError: If 'kernel_size' is not a 2-D list
  """
  assert len(tensors_list) > 1, "Merge required 2 or more tensors."

  with tf.name_scope(name):
    tensors = [l for l in tensors_list]
    if mode == 'concat':
      output = tf.concat(tensors, axis=axis)
    elif mode == 'elemwise_sum':
      output = tensors[0]
      for i in range(1, len(tensors)):
        output = tf.add(output, tensors[i])
    elif mode == 'elemwise_mul':
      output = tensors[0]
      for i in range(1, len(tensors)):
        output = tf.multiply(output, tensors[i])
    elif mode == 'sum':
      output = tf.reduce_sum(tf.concat(tensors, axis=axis), axis=axis)
    elif mode == 'mean':
      output = tf.reduce_mean(tf.concat(tensors, axis=axis), axis=axis)
    elif mode == 'prod':
      output = tf.reduce_prod(tf.concat(tensors, axis=axis), axis=axis)
    elif mode == 'max':
      output = tf.reduce_max(tf.concat(tensors, axis=axis), axis=axis)
    elif mode == 'min':
      output = tf.reduce_min(tf.concat(tensors, axis=axis), axis=axis)
    elif mode == 'and':
      output = tf.reduce_all(tf.concat(tensors, axis=axis), axis=axis)
    elif mode == 'or':
      output = tf.reduce_any(tf.concat(tensors, axis=axis), axis=axis)
    else:
      raise Exception("Unknown merge mode", str(mode))
    return _collect_named_outputs(outputs_collections, name, output)

  return output


def stack(inputs, layer, stack_args, is_training, reuse, outputs_collections=None, **kwargs):
  """Builds a stack of layers by applying layer repeatedly using stack_args.
    `stack` allows you to repeatedly apply the same operation with different
    arguments `stack_args[i]`. For each application of the layer, `stack` creates
    a new scope appended with an increasing number. For example:
    ```python
        y = stack(x, fully_connected, [32, 64, 128], scope='fc')
       # It is equivalent to:
       x = fully_connected(x, 32, scope='fc/fc_1')
       x = fully_connected(x, 64, scope='fc/fc_2')
       y = fully_connected(x, 128, scope='fc/fc_3')
    ```
    If the `scope` argument is not given in `kwargs`, it is set to
    `layer.__name__`, or `layer.func.__name__` (for `functools.partial`
    objects). If neither `__name__` nor `func.__name__` is available, the
    layers are called with `scope='stack'`.

    Args:
        inputs: A `Tensor` suitable for layer.
        layer: A layer with arguments `(inputs, *args, **kwargs)`
        stack_args: A list/tuple of parameters for each call of layer.
        outputs_collections: The collections to which the outputs are added.
        **kwargs: Extra kwargs for the layer.

    Returns:
       a `Tensor` result of applying the stacked layers.

    Raises:
        ValueError: if the op is unknown or wrong.
    """
  name = kwargs.pop('name', None)
  if not isinstance(stack_args, (list, tuple)):
    raise ValueError('stack_args need to be a list or tuple')
  with tf.variable_scope(name, 'Stack', [inputs]):
    inputs = tf.convert_to_tensor(inputs)
    if name is None:
      if hasattr(layer, '__name__'):
        name = layer.__name__
      elif hasattr(layer, 'func') and hasattr(layer.func, '__name__'):
        # In case layer is a functools.partial.
        name = layer.func.__name__
      else:
        name = 'stack'
    outputs = inputs
    for i in range(len(stack_args)):
      kwargs['name'] = name + '_' + str(i + 1)
      layer_args = stack_args[i]
      # if not isinstance(layer_args, (list, tuple)):
      #    layer_args = [layer_args]
      outputs = layer(outputs, layer_args, is_training, reuse, **kwargs)
    return _collect_named_outputs(outputs_collections, name, outputs)


def unit_norm(inputs, dim, epsilon=1e-7, scope=None):
  """Normalizes the given input across the specified dimension to unit length.
  Note that the rank of `input` must be known.

  Args:
      inputs: A `Tensor` of arbitrary size.
      dim: The dimension along which the input is normalized.
      epsilon: A small value to add to the inputs to avoid dividing by zero.
      scope: Optional scope for variable_scope.

  Returns:
      The normalized `Tensor`.

  Raises:
      ValueError: If dim is larger than the number of dimensions in 'inputs'.
  """
  with tf.variable_scope(scope, 'UnitNorm', [inputs]):
    if not inputs.get_shape():
      raise ValueError('The input rank must be known.')
    input_rank = len(inputs.get_shape().as_list())
    if dim < 0 or dim >= input_rank:
      raise ValueError('dim must be positive but smaller than the input rank.')

    lengths = tf.sqrt(epsilon + tf.reduce_sum(tf.square(inputs), dim, True))
    multiples = []
    if dim > 0:
      multiples.append(tf.ones([dim], tf.int32))
    multiples.append(tf.strided_slice(tf.shape(inputs), [dim], [dim + 1], [1]))
    if dim < (input_rank - 1):
      multiples.append(tf.ones([input_rank - 1 - dim], tf.int32))
    multiples = tf.concat(multiples, 0)
    return tf.div(inputs, tf.tile(lengths, multiples))


def crop_and_concat(inputs1, inputs2, name='crop_concat'):
  """Concates two features maps concates different sizes feature maps cropping
  the larger map concatenation across output channels.

  Args:
      inputs1: A `Tensor`
      inputs2: A `Tensor`

  Returns:
     concated output tensor
  """
  with tf.name_scope(name):
    inputs1_shape = tf.shape(inputs1)
    inputs2_shape = tf.shape(inputs2)
    # offsets for the top left corner of the crop
    offsets = [
        0, (inputs1_shape[1] - inputs2_shape[1]) // 2, (inputs1_shape[2] - inputs2_shape[2]) // 2, 0
    ]
    size = [-1, inputs2_shape[1], inputs2_shape[2], -1]
    inputs1_crop = tf.slice(inputs1, offsets, size)
    return tf.concat([inputs1_crop, inputs2], axis=3)


class GradientReverseLayer(object):

  def __init__(self):
    self.num_calls = 0

  def __call__(self, x, gamma=1.0):
    grad_name = "GradientReverse%d" % self.num_calls

    @ops.RegisterGradient(grad_name)
    def _gradients_reverse(op, grad):
      return [tf.neg(grad) * gamma]

    g = tf.get_default_graph()
    with g.gradient_override_map({"Identity": grad_name}):
      y = tf.identity(x)

    self.num_calls += 1
    return y


gradient_reverse = GradientReverseLayer()


def _collect_named_outputs(outputs_collections, name, output):
  if outputs_collections is not None:
    tf.add_to_collection(outputs_collections, NamedOutputs(name, output))
  return output


def _check_unused(unused, name):
  allowed_keys = ['is_training', 'reuse', 'outputs_collections', 'trainable', 'padding']
  helper.veryify_args(unused, allowed_keys, 'Layer "%s" got unexpected argument(s):' % name)
