text: import numpy as np
  import tensorflow as tf
  from tensorflow.python.framework import ops
  import pydensecrf.densecrf as dcrf
  from .layers import conv2d, depthwise_conv2d, avg_pool_2d, \
      max_pool, relu, crelu, batch_norm_tf as batch_norm
  from ..utils import util
  from . import initializers as initz
  
  
  def spatialtransformer(U,
                         theta,
                         batch_size=64,
                         downsample_factor=1.0,
                         num_transform=1,
                         name='SpatialTransformer',
                         **kwargs):
    """Spatial Transformer Layer.
  
    Implements a spatial transformer layer as described in [1]_.
    It's based on lasagne implementation in [2]_, modified by Mrinal Haloi
  
    Args:
        U: float
            The output of a convolutional net should have the
            shape [batch_size, height, width, num_channels].
        theta: float
            The output of the localisation network should be [batch_size,
            num_transform, 6] or [batch_size, 6] if num_transform=1
            ```python
                `theta`` to :
                    identity = np.array([[1., 0., 0.],
                                     [0., 1., 0.]])
                    identity = identity.flatten()
                    theta = tf.Variable(initial_value=identity)
            ```
        downsample_factor: a float, determines output shape, downsample input
             shape by downsample_factor
  
    Returns:
        spatial transformed output of the network
  
    References
    .. [1]  "Spatial Transformer Networks", Max Jaderberg, Karen Simonyan,
        Andrew Zisserman, Koray Kavukcuoglu
    .. [2]  https://github.com/skaae/transformer_network/blob/master/transformerlayer.py
    """
    with tf.variable_scope(name):
      if num_transform > 1 and len(theta.get_shape().as_list()) == 3:
        _, num_transforms = map(int, theta.get_shape().as_list()[:2])
        indices = [[i] * num_transforms for i in range(batch_size)]
        U = tf.gather(U, tf.reshape(indices, [-1]))
  
      input_shape = U.get_shape().as_list()
      num_channels = input_shape[3]
      theta = tf.reshape(theta, (-1, 2, 3))
      theta = tf.cast(theta, tf.float32)
      if not isinstance(downsample_factor, float):
        downsample_factor = tf.cast(downsample_factor, tf.float32)
  
      # grid of (x_t, y_t, 1), eq (1) in ref [1]
      out_height = tf.cast(input_shape[1] / downsample_factor, tf.int32)
      out_width = tf.cast(input_shape[2] / downsample_factor, tf.int32)
      grid = _meshgrid(out_height, out_width)
      grid = tf.expand_dims(grid, 0)
      grid = tf.reshape(grid, [-1])
      grid = tf.tile(grid, tf.stack([batch_size]))
      grid = tf.reshape(grid, tf.stack([batch_size, 3, -1]))
  
      # Transform A x (x_t, y_t, 1)^T -> (x_s, y_s)
      T_g = tf.matmul(theta, grid)
      x_s = tf.slice(T_g, [0, 0, 0], [-1, 1, -1])
      y_s = tf.slice(T_g, [0, 1, 0], [-1, 1, -1])
      x_s_flat = tf.reshape(x_s, [-1])
      y_s_flat = tf.reshape(y_s, [-1])
  
      input_transformed = _interpolate(U, x_s_flat, y_s_flat, batch_size, downsample_factor)
  
      output = tf.reshape(input_transformed,
                          tf.stack([batch_size, out_height, out_width, num_channels]))
      return output
  
  
  def _repeat(x, n_repeats):
    with tf.variable_scope('_repeat'):
      rep = tf.transpose(tf.expand_dims(tf.ones(shape=tf.stack([
          n_repeats,
      ])), 1), [1, 0])
      rep = tf.cast(rep, tf.int32)
      x = tf.matmul(tf.reshape(x, (-1, 1)), rep)
      return tf.reshape(x, [-1])
  
  
  def _interpolate(im, x, y, batch_size, downsample_factor):
    with tf.variable_scope('_interpolate'):
      input_shape = im.get_shape().as_list()
      height = input_shape[1]
      width = input_shape[2]
      channels = input_shape[3]
  
      x = tf.cast(x, tf.float32)
      y = tf.cast(y, tf.float32)
      height_f = tf.cast(height, tf.float32)
      width_f = tf.cast(width, tf.float32)
      out_height = tf.cast(height / downsample_factor, tf.int32)
      out_width = tf.cast(width / downsample_factor, tf.int32)
      zero = tf.zeros([], dtype=tf.int32)
      max_y = tf.cast(height - 1, tf.int32)
      max_x = tf.cast(width - 1, tf.int32)
  
      # scale indices from [-1, 1] to [0, width/height]
      x = (x + 1.0) * (width_f) / 2.0
      y = (y + 1.0) * (height_f) / 2.0
  
      # do sampling
      x0 = tf.cast(tf.floor(x), tf.int32)
      x1 = x0 + 1
      y0 = tf.cast(tf.floor(y), tf.int32)
      y1 = y0 + 1
  
      x0 = tf.clip_by_value(x0, zero, max_x)
      x1 = tf.clip_by_value(x1, zero, max_x)
      y0 = tf.clip_by_value(y0, zero, max_y)
      y1 = tf.clip_by_value(y1, zero, max_y)
      dim2 = width
      dim1 = width * height
      base = _repeat(tf.range(batch_size) * dim1, out_height * out_width)
      base_y0 = base + y0 * dim2
      base_y1 = base + y1 * dim2
      idx_a = base_y0 + x0
      idx_b = base_y1 + x0
      idx_c = base_y0 + x1
      idx_d = base_y1 + x1
  
      # use indices to lookup pixels in the flat image and restore
      # channels dim
      im_flat = tf.reshape(im, tf.stack([-1, channels]))
      im_flat = tf.cast(im_flat, tf.float32)
      Ia = tf.gather(im_flat, idx_a)
      Ib = tf.gather(im_flat, idx_b)
      Ic = tf.gather(im_flat, idx_c)
      Id = tf.gather(im_flat, idx_d)
  
      # and finally calculate interpolated values
      x0_f = tf.cast(x0, tf.float32)
      x1_f = tf.cast(x1, tf.float32)
      y0_f = tf.cast(y0, tf.float32)
      y1_f = tf.cast(y1, tf.float32)
      wa = tf.expand_dims(((x1_f - x) * (y1_f - y)), 1)
      wb = tf.expand_dims(((x1_f - x) * (y - y0_f)), 1)
      wc = tf.expand_dims(((x - x0_f) * (y1_f - y)), 1)
      wd = tf.expand_dims(((x - x0_f) * (y - y0_f)), 1)
      output = tf.add_n([wa * Ia, wb * Ib, wc * Ic, wd * Id])
      return output
  
  
  def _meshgrid(height, width):
    with tf.variable_scope('_meshgrid'):
      x_t = tf.matmul(
          tf.ones(shape=tf.stack([height, 1])),
          tf.transpose(tf.expand_dims(tf.linspace(-1.0, 1.0, width), 1), [1, 0]))
      y_t = tf.matmul(
          tf.expand_dims(tf.linspace(-1.0, 1.0, height), 1), tf.ones(shape=tf.stack([1, width])))
  
      x_t_flat = tf.reshape(x_t, (1, -1))
      y_t_flat = tf.reshape(y_t, (1, -1))
  
      ones = tf.ones_like(x_t_flat)
      grid = tf.concat([x_t_flat, y_t_flat, ones], 0)
      return grid
  
  
  def subsample(inputs, factor, name=None):
    """Subsamples the input along the spatial dimensions.
  
    Args:
        inputs: A `Tensor` of size [batch, height_in, width_in, channels].
        factor: The subsampling factor.
        name: Optional variable_scope.
  
    Returns:
        output: A `Tensor` of size [batch, height_out, width_out, channels] with the
            input, either intact (if factor == 1) or subsampled (if factor > 1).
    """
    if factor == 1:
      return inputs
    else:
      return max_pool(inputs, filter_size=(1, 1), stride=(factor, factor), name=name)
  
  
  def conv2d_same(inputs, num_outputs, kernel_size, stride, name=None, **kwargs):
    """Strided 2-D convolution with 'SAME' padding.
  
    When stride > 1, then we do explicit zero-padding, followed by conv2d with
    'VALID' padding.
  
    Note that
  
       net = conv2d_same(inputs, num_outputs, 3, stride=stride)
  
    is equivalent to
  
       net = conv2d(inputs, num_outputs, 3, stride=1, padding='SAME')
       net = subsample(net, factor=stride)
  
    whereas
  
       net = conv2d(inputs, num_outputs, 3, stride=stride, padding='SAME')
  
    is different when the input's height or width is even, which is why we add the
    current function. For more details, see ResnetUtilsTest.testConv2DSameEven().
  
    Args:
        inputs: A 4-D tensor of size [batch, height_in, width_in, channels].
        num_outputs: An integer, the number of output filters.
        kernel_size: An int with the kernel_size of the filters.
        stride: An integer, the output stride.
        name: name.
  
    Returns:
        output: A 4-D tensor of size [batch, height_out, width_out, channels] with
            the convolution output.
    """
    if stride == 1:
      return conv2d(
          inputs,
          num_outputs,
          filter_size=(kernel_size, kernel_size),
          stride=(1, 1),
          padding='SAME',
          name=name,
          **kwargs)
    else:
      kernel_size_effective = kernel_size + (kernel_size - 1) * (rate - 1)
      pad_total = kernel_size_effective - 1
      pad_beg = pad_total // 2
      pad_end = pad_total - pad_beg
      inputs = tf.pad(inputs, [[0, 0], [pad_beg, pad_end], [pad_beg, pad_end], [0, 0]])
      return conv2d(
          inputs, num_outputs, kernel_size, stride=stride, padding='VALID', name=name, **kwargs)
  
  
  def bottleneck_v1(inputs, depth, depth_bottleneck, stride, rate=1, name=None, **kwargs):
    """Bottleneck residual unit variant with BN before convolutions.
  
    This is the full preactivation residual unit variant proposed in [2]. See
    Fig. 1(b) of [2] for its definition. Note that we use here the bottleneck
    variant which has an extra bottleneck layer.
  
    When putting together two consecutive ResNet blocks that use this unit, one
    should use stride = 2 in the last unit of the first block.
  
    Args:
        inputs: A tensor of size [batch, height, width, channels].
        depth: The depth of the ResNet unit output.
        depth_bottleneck: The depth of the bottleneck layers.
        stride: The ResNet unit's stride. Determines the amount of downsampling of
            the units output compared to its input.
        rate: An integer, rate for atrous convolution.
        outputs_collections: Collection to add the ResNet unit output.
        name: Optional variable_scope.
  
    Returns:
        The ResNet unit's output.
    """
    is_training = kwargs.get('is_training')
    reuse = kwargs.get('reuse')
    with tf.variable_scope(name, 'bottleneck_v2', [inputs]):
      depth_in = util.last_dimension(inputs.get_shape(), min_rank=4)
      preact = batch_norm(
          inputs, activation_fn=tf.nn.relu, name='preact', is_training=is_training, reuse=reuse)
      if depth == depth_in:
        shortcut = subsample(inputs, stride, 'shortcut')
      else:
        shortcut = conv2d(
            preact,
            depth,
            is_training,
            reuse,
            filter_size=(1, 1),
            stride=(stride, stride),
            batch_norm=None,
            activation=None,
            name='shortcut')
  
      residual = conv2d(
          preact, depth_bottleneck, filter_size=(1, 1), stride=(1, 1), name='conv1', **kwargs)
      residual = conv2d_same(residual, depth_bottleneck, 3, stride, rate=rate, name='conv2', **kwargs)
      residual = conv2d(
          residual,
          depth,
          is_training,
          reuse,
          filter_size=(1, 1),
          stride=(1, 1),
          batch_norm=None,
          activation=None,
          name='conv3')
  
      output = shortcut + residual
  
      return output
  
  
  def bottleneck_v2(inputs, depth, depth_bottleneck, stride, rate=1, name=None, **kwargs):
    """Bottleneck residual unit variant with BN before convolutions.
  
    This is the full preactivation residual unit variant proposed in [2]. See
    Fig. 1(b) of [2] for its definition. Note that we use here the bottleneck
    variant which has an extra bottleneck layer.
  
    When putting together two consecutive ResNet blocks that use this unit, one
    should use stride = 2 in the last unit of the first block.
  
    Args:
        inputs: A tensor of size [batch, height, width, channels].
        depth: The depth of the ResNet unit output.
        depth_bottleneck: The depth of the bottleneck layers.
        stride: The ResNet unit's stride. Determines the amount of downsampling of
            the units output compared to its input.
        rate: An integer, rate for atrous convolution.
        outputs_collections: Collection to add the ResNet unit output.
        name: Optional variable_scope.
  
    Returns:
        The ResNet unit's output.
    """
    is_training = kwargs.get('is_training')
    reuse = kwargs.get('reuse')
    with tf.variable_scope(name, 'bottleneck_v2', [inputs]):
      depth_in = util.last_dimension(inputs.get_shape(), min_rank=4)
      preact = batch_norm(
          inputs, activation_fn=tf.nn.relu, name='preact', is_training=is_training, reuse=reuse)
      if depth == depth_in:
        shortcut = subsample(inputs, stride, 'shortcut')
      else:
        shortcut = conv2d(
            preact,
            depth,
            is_training,
            reuse,
            filter_size=(1, 1),
            stride=(stride, stride),
            batch_norm=None,
            activation=None,
            name='shortcut')
  
      residual = conv2d(
          preact, depth_bottleneck, filter_size=(1, 1), stride=(1, 1), name='conv1', **kwargs)
      residual = conv2d_same(residual, depth_bottleneck, 3, stride, rate=rate, name='conv2', **kwargs)
      residual = conv2d(
          residual,
          depth,
          is_training,
          reuse,
          filter_size=(1, 1),
          stride=(1, 1),
          batch_norm=None,
          activation=None,
          name='conv3')
  
      output = tf.nn.relu(shortcut + residual)
  
      return output
  
  
  def memory_module(inputs,
                    time,
                    context,
                    reuse,
                    nwords,
                    edim,
                    mem_size,
                    lindim,
                    batch_size,
                    nhop,
                    share_list,
                    regularizer=tf.nn.l2_loss,
                    init_f=initz.random_normal(),
                    trainable=True,
                    name=None,
                    **kwargs):
    with tf.variable_scope(name, 'memory_module', reuse=reuse):
      A_shape = B_shape = [nwords, edim]
      C_shape = [edim, edim]
      T_A_shape = T_B_shape = [mem_size, edim]
      A = tf.get_variable(
          name='A', shape=A_shape, initializer=init_f, regularizer=regularizer, trainable=trainable)
      B = tf.get_variable(
          name='B', shape=B_shape, initializer=init_f, regularizer=regularizer, trainable=trainable)
      C = tf.get_variable(
          name='C', shape=C_shape, initializer=init_f, regularizer=regularizer, trainable=trainable)
      T_A = tf.get_variable(
          name='T_A',
          shape=T_A_shape,
          initializer=init_f,
          regularizer=regularizer,
          trainable=trainable)
      T_B = tf.get_variable(
          name='T_B',
          shape=T_B_shape,
          initializer=init_f,
          regularizer=regularizer,
          trainable=trainable)
  
      # m_i = sum A_ij * x_ij + T_A_i
      Ain_c = tf.nn.embedding_lookup(A, context)
      Ain_t = tf.nn.embedding_lookup(T_A, time)
      Ain = tf.add(Ain_c, Ain_t)
  
      # c_i = sum B_ij * u + T_B_i
      Bin_c = tf.nn.embedding_lookup(B, context)
      Bin_t = tf.nn.embedding_lookup(T_B, time)
      Bin = tf.add(Bin_c, Bin_t)
  
      hid = []
      hid.append(inputs)
  
      for h in range(nhop):
        hid3dim = tf.reshape(hid[-1], [-1, 1, edim])
        Aout = tf.matmul(hid3dim, Ain, adjoint_b=True)
        Aout2dim = tf.reshape(Aout, [-1, mem_size])
        P = tf.nn.softmax(Aout2dim)
  
        probs3dim = tf.reshape(P, [-1, 1, mem_size])
        Bout = tf.matmul(probs3dim, Bin)
        Bout2dim = tf.reshape(Bout, [-1, edim])
  
        Cout = tf.matmul(hid[-1], C)
        Dout = tf.add(Cout, Bout2dim)
  
        share_list[0].append(Cout)
  
        if lindim == edim:
          hid.append(Dout)
        elif lindim == 0:
          hid.append(tf.nn.relu(Dout))
        else:
          F = tf.slice(Dout, [0, 0], [batch_size, lindim])
          G = tf.slice(Dout, [0, lindim], [batch_size, edim - lindim])
          K = tf.nn.relu(G)
          hid.append(tf.concat([F, K], 1))
      return hid
  
  
  def dense_crf(probs,
                img=None,
                n_classes=21,
                n_iters=10,
                sxy_gaussian=(1, 1),
                compat_gaussian=4,
                kernel_gaussian=dcrf.DIAG_KERNEL,
                normalisation_gaussian=dcrf.NORMALIZE_SYMMETRIC,
                sxy_bilateral=(49, 49),
                compat_bilateral=2,
                srgb_bilateral=(13, 13, 13),
                kernel_bilateral=dcrf.DIAG_KERNEL,
                normalisation_bilateral=dcrf.NORMALIZE_SYMMETRIC):
    """DenseCRF over unnormalised predictions. More details on the arguments at
    https://github.com/lucasb-eyer/pydensecrf.
  
    Args:
        probs: class probabilities per pixel.
        img: if given, the pairwise bilateral potential on raw RGB values will be computed.
        n_iters: number of iterations of MAP inference.
        sxy_gaussian: standard deviations for the location component
            of the colour-independent term.
        compat_gaussian: label compatibilities for the colour-independent
            term (can be a number, a 1D array, or a 2D array).
        kernel_gaussian: kernel precision matrix for the colour-independent
            term (can take values CONST_KERNEL, DIAG_KERNEL, or FULL_KERNEL).
        normalisation_gaussian: normalisation for the colour-independent term
            (possible values are NO_NORMALIZATION, NORMALIZE_BEFORE, NORMALIZE_AFTER,
             NORMALIZE_SYMMETRIC).
        sxy_bilateral: standard deviations for the location component of the colour-dependent term.
        compat_bilateral: label compatibilities for the colour-dependent
            term (can be a number, a 1D array, or a 2D array).
        srgb_bilateral: standard deviations for the colour component
            of the colour-dependent term.
        kernel_bilateral: kernel precision matrix for the colour-dependent term
            (can take values CONST_KERNEL, DIAG_KERNEL, or FULL_KERNEL).
        normalisation_bilateral: normalisation for the colour-dependent term
            (possible values are NO_NORMALIZATION, NORMALIZE_BEFORE, NORMALIZE_AFTER,
            NORMALIZE_SYMMETRIC).
  
    Returns:
        Refined predictions after MAP inference.
    """
    _, h, w, _ = probs.shape
  
    probs = probs[0].transpose(2, 0, 1).copy(order='C')
  
    d = dcrf.DenseCRF2D(h, w, n_classes)
    U = -np.log(probs)
    U = U.reshape((n_classes, -1))
    d.setUnaryEnergy(U)
    d.addPairwiseGaussian(
        sxy=sxy_gaussian,
        compat=compat_gaussian,
        kernel=kernel_gaussian,
        normalization=normalisation_gaussian)
    if img is not None:
      assert (img.shape[1:3] == (
          h, w)), "The image height and width must coincide with dimensions of the logits."
      d.addPairwiseBilateral(
          sxy=sxy_bilateral,
          compat=compat_bilateral,
          kernel=kernel_bilateral,
          normalization=normalisation_bilateral,
          srgb=srgb_bilateral,
          rgbim=img[0])
    Q = d.inference(n_iters)
    preds = np.array(Q, dtype=np.float32).reshape((n_classes, h, w)).transpose(1, 2, 0)
    preds = np.expand_dims(preds, 0)
    return preds
  
  
  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.negative(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 resnext_block(inputs,
                    nb_blocks,
                    out_channels,
                    is_training,
                    reuse,
                    cardinality,
                    downsample=False,
                    downsample_strides=2,
                    activation=relu,
                    batch_norm=None,
                    batch_norm_args=None,
                    name="ResNeXtBlock",
                    **kwargs):
    """ResNeXt Block. resnext paper https://arxiv.org/pdf/1611.05431.pdf.
  
    Args:
        inputs: `Tensor`. Inputs 4-D Layer.
        nb_blocks: `int`. Number of layer blocks.
        out_channels: `int`. The number of convolutional filters of the
            layers surrounding the bottleneck layer.
        cardinality: `int`. Number of aggregated residual transformations.
        downsample: `bool`. If True, apply downsampling using
            'downsample_strides' for strides.
        downsample_strides: `int`. The strides to use when downsampling.
        activation: `function` (returning a `Tensor`).
        batch_norm: `bool`. If True, apply batch normalization.
        use_ bias: `bool`. If True, a bias is used.
        w_init: `function`, Weights initialization.
        b_init: `tf.Tensor`. Bias initialization.
        w_regularizer: `function`. Add a regularizer to this
        weight_decay: `float`. Regularizer decay parameter. Default: 0.001.
        trainable: `bool`. If True, weights will be trainable.
        reuse: `bool`. If True and 'scope' is provided, this layer variables
            will be reused (shared).
            override name.
        name: A name for this layer (optional). Default: 'ResNeXtBlock'.
  
    Returns:
        4-D Tensor [batch, new height, new width, out_channels].
    """
    resnext = inputs
    input_shape = util.get_input_shape(inputs)
    in_channels = input_shape[-1]
  
    card_values = [1, 2, 4, 8, 32]
    bottleneck_values = [64, 40, 24, 14, 4]
    bottleneck_size = bottleneck_values[card_values.index(cardinality)]
    group_width = [64, 80, 96, 112, 128]
  
    assert cardinality in card_values, "cardinality must be in [1, 2, 4, 8, 32]"
  
    with tf.variable_scope(name, reuse=reuse):
      for i in range(nb_blocks):
        identity = resnext
        if not downsample:
          downsample_strides = 1
  
        resnext = conv2d(
            resnext,
            bottleneck_size,
            is_training,
            reuse,
            filter_size=1,
            stride=downsample_strides,
            batch_norm=batch_norm,
            batch_norm_args=batch_norm_args,
            activation=activation,
            padding='valid',
            name='conv1',
            **kwargs)
  
        resnext = depthwise_conv2d(
            resnext,
            cardinality,
            is_training,
            reuse,
            filter_size=3,
            batch_norm=batch_norm,
            batch_norm_args=batch_norm_args,
            activation=relu,
            stride=1,
            **kwargs)
        resnext = conv2d(
            resnext,
            out_channels,
            is_training,
            reuse,
            filter_size=1,
            stride=1,
            batch_norm=None,
            activation=activation,
            padding='valid',
            name='conv2',
            **kwargs)
  
        if batch_norm is not None:
          batch_norm_args = batch_norm_args or {}
          resnext = batch_norm(
              resnext, is_training=is_training, reuse=reuse, trainable=trainable, **batch_norm_args)
        if downsample_strides > 1:
          identity = avg_pool_2d(identity, filter_size=1, stride=downsample_strides)
  
        if in_channels != out_channels:
          ch = (out_channels - in_channels) // 2
          identity = tf.pad(identity, [[0, 0], [0, 0], [0, 0], [ch, ch]])
          in_channels = out_channels
  
        resnext = resnext + identity
        resnext = activation(resnext)
  
      return resnext
  
  
  def embedding(inputs,
                vocab_dim,
                embedding_dim,
                reuse,
                validate_indices=False,
                w_init=tf.random_uniform_initializer(-1., 1.),
                trainable=True,
                normalize=False,
                vocab_freqs=None,
                name="Embedding"):
    """Embedding. Embedding layer for a sequence of integer ids or floats.
  
    Args:
        inputs: a 2-D `Tensor` [samples, ids].
        vocab_dim: list of `int`. Vocabulary size (number of ids).
        embedding_dim: list of `int`. Embedding size.
        validate_indices: `bool`. Whether or not to validate gather indices.
        w_init:  Weights initialization.
        trainable: `bool`. If True, weights will be trainable.
        reuse: `bool`. If True and 'scope' is provided, this layer variables
            will be reused (shared).
        name: A name for this layer (optional). Default: 'Embedding'.
  
    Returns:
        3-D Tensor [samples, embedded_ids, features].
    """
  
    input_shape = util.get_input_shape(inputs)
    assert len(input_shape) == 2, "Input Tensor shape must be 2-D"
  
    with tf.variable_scope(name, reuse=reuse):
      with tf.device('/cpu:0'):
        W = tf.get_variable(
            "W", shape=[vocab_dim, embedding_dim], initializer=w_init, trainable=trainable)
      if normalize:
        assert vocab_freqs is not None
        vocab_freqs = tf.constant(vocab_freqs, dtype=tf.float32, shape=(vocab_dim, 1))
        W = _normalize(W, vocab_freqs)
  
      output = tf.cast(inputs, tf.int32)
      output = tf.nn.embedding_lookup(W, output, validate_indices=validate_indices)
  
    shape = [-1] + output.get_shape().as_list()[1:3] + [1]
    # seq_length = util.retrieve_seq_length(tf.reshape(inputs, shape))
  
    return output
  
  
  def _normalize(emb, vocab_freqs):
    weights = vocab_freqs / tf.reduce_sum(vocab_freqs)
    mean = tf.reduce_sum(weights * emb, 0, keep_dims=True)
    var = tf.reduce_sum(weights * tf.pow(emb - mean, 2.), 0, keep_dims=True)
    stddev = tf.sqrt(1e-6 + var)
    return (emb - mean) / stddev
  
  
  def gated_layer(inputs, layer, num_units, is_training, reuse, name='gated_layer', **kwargs):
    """Gated unit for language modelling.
  
    Args:
        inputs: a 3-D/4-D `Tensor`, input [samples, timesteps, input_dim]
        layer: a `layer`, layer to pass the inputs
            e.g. `tefla.core.layers`
        num_units: a `int`, number of units for each layer
        is_training: a `boolean`, Training if its true
        reuse: `bool`. If True and 'scope' is provided, this layer variables
            will be reused (shared).
        name: A name for this layer (optional). Default: 'gated_layer'.
  
    Returns:
        a 3-D/4-D `Tensor`, output of the gated unit
    """
    activation = kwargs.get('activation')
    with tf.name_scope(name):
      inputs = activation(inputs)
      inputs_shortcut = layer(
          inputs, num_units, is_training, reuse, filter_size=1, name=name + '_shortcut', **kwargs)
      output1 = layer(inputs, num_units, is_training, reuse, name=name + '_layer1', **kwargs)
      output2 = layer(inputs, num_units, is_training, reuse, name=name + '_layer2', **kwargs)
      output = tf.multiply(output1, tf.sigmoid(output2))
      return output + inputs_shortcut
  
  
  def glimpseSensor(img, normLoc, minRadius=4, depth=1, sensorBandwidth=12):
    """Returns glimpses at the locations.
  
    Args:
        img: a 4-D `Tensor`, [batch_size, width, height, channels]
        normloc: a `float`, [0, 1] normalized location
        minRadius: a `int`, min radius for zooming
        depth: a `int`, number of zooms
        sensorbandwidth: a `int`, output glimpse size, width/height
  
    Returns:
        a 5-D `tensor` of glimpses
    """
    input_shape = img.get_shape().as_list()
    img_size = input_shape[1]
    loc = tf.round(((normLoc + 1) / 2.0) * img_size)
    loc = tf.cast(loc, tf.int32)
    zooms = []
    for k in range(batch_size):
      imgZooms = []
      one_img = img[k, :, :, :]
      max_radius = minRadius * (2**(depth - 1))
      offset = 2 * max_radius
      one_img = tf.image.pad_to_bounding_box(one_img, offset, offset, max_radius * 4 + img_size,
                                             max_radius * 4 + img_size)
  
      for i in range(depth):
        r = int(minRadius * (2**(i)))
        d_raw = 2 * r
        d = tf.constant(d_raw, shape=[1])
        d = tf.tile(d, [2])
        loc_k = loc[k, :]
        adjusted_loc = offset + loc_k - r
        one_img2 = tf.reshape(
            one_img, (one_img.get_shape()[0].value, one_img.get_shape()[1].value, input_shape[3]))
        zoom = tf.slice(one_img2, adjusted_loc, d)
        zoom = tf.image.resize_bilinear(
            tf.reshape(zoom, (1, d_raw, d_raw, nput_shape[3])), (sensorBandwidth, sensorBandwidth))
        zoom = tf.reshape(zoom, (sensorBandwidth, sensorBandwidth, input_shape[3]))
        imgZooms.append(zoom)
  
      zooms.append(tf.stack(imgZooms))
  
    return tf.stack(zooms)
  
  
  def pva_block_v1(x, num_units, name='pva_block_v1', **kwargs):
    """Adds a PVA block layer. convolution followed by crelu and scaling.
  
    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
        num_units: 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.
        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.
    """
    trainable = kwargs.get('trainable')
    regularizer = kwargs.get('w_regularizer')
    activation = kwargs.get('activation')
    kwargs.update({'activation': None})
    with tf.variable_scope(name):
      output = conv2d(x, num_units, **kwargs)
      output = crelu(output)
      alpha = tf.get_variable(
          'scale/alpha',
          shape=[
              2 * num_units,
          ],
          dtype=tf.float32,
          initializer=tf.constant_initializer(1.0),
          trainable=trainable,
          regularizer=regularizer)
      beta = tf.get_variable(
          'scale/beta',
          shape=[
              2 * num_units,
          ],
          dtype=tf.float32,
          initializer=tf.constant_initializer(0.0),
          trainable=trainable,
          regularizer=regularizer)
      output = tf.add(tf.multiply(output, alpha), beta)
      return activation(output)
  
  
  def pva_block_v2(x, num_units, name='pva_block_v2', **kwargs):
    """Adds a PVA block v2 layer. first batch normalization followed by crelu and
    scaling, convolution is applied after scalling.
  
    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
        num_units: 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.
        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.
    """
    trainable = kwargs.get('trainable')
    reuse = kwargs.get('reuse')
    is_training = kwargs.get('is_training')
    regularizer = kwargs.get('w_regularizer')
    activation = kwargs.get('activation')
    kwargs.update({'activation': None})
    kwargs.update({'batch_norm': None})
    with tf.variable_scope(name, reuse=reuse):
      output = batch_norm(
          x, activation_fn=tf.nn.relu, name='batch_norm', is_training=is_training, reuse=reuse)
      output = crelu(output)
      alpha = tf.get_variable(
          'scale/alpha',
          shape=[
              2 * num_units,
          ],
          dtype=tf.float32,
          initializer=tf.constant_initializer(1.0),
          trainable=trainable,
          regularizer=regularizer)
      beta = tf.get_variable(
          'scale/beta',
          shape=[
              2 * num_units,
          ],
          dtype=tf.float32,
          initializer=tf.constant_initializer(0.0),
          trainable=trainable,
          regularizer=regularizer)
      output = tf.add(tf.multiply(output, alpha), beta)
      output = activation(output)
      output = conv2d(x, num_units, **kwargs)
      return output
  
  
  def max_pool_2d_nxn_regions(inputs, output_size, mode='max'):
    """Performs a pooling operation that results in a fixed size: output_size x
    output_size.
  
    Used by spatial_pyramid_pool. Refer to appendix A in [1].
  
    Args:
        inputs: A 4D Tensor (B, H, W, C)
        output_size: The output size of the pooling operation.
        mode: The pooling mode {max, avg}
  
    Returns:
        A list of tensors, for each output bin.
        The list contains output_size * output_size elements, where
        each elment is a Tensor (N, C).
    """
    inputs_shape = tf.shape(inputs)
    h = tf.cast(tf.gather(inputs_shape, 1), tf.int32)
    w = tf.cast(tf.gather(inputs_shape, 2), tf.int32)
  
    if mode == 'max':
      pooling_op = tf.reduce_max
    elif mode == 'avg':
      pooling_op = tf.reduce_mean
    else:
      msg = "Mode must be either 'max' or 'avg'. Got '{0}'"
      raise ValueError(msg.format(mode))
  
    result = []
    n = output_size
    for row in range(output_size):
      for col in range(output_size):
        # start_h = floor(row / n * h)
        start_h = tf.cast(tf.floor(tf.mul(tf.divide(row, n), tf.cast(h, tf.float32))), tf.int32)
        # end_h = ceil((row + 1) / n * h)
        end_h = tf.cast(tf.ceil(tf.mul(tf.divide((row + 1), n), tf.cast(h, tf.float32))), tf.int32)
        # start_w = floor(col / n * w)
        start_w = tf.cast(tf.floor(tf.mul(tf.divide(col, n), tf.cast(w, tf.float32))), tf.int32)
        # end_w = ceil((col + 1) / n * w)
        end_w = tf.cast(tf.ceil(tf.mul(tf.divide((col + 1), n), tf.cast(w, tf.float32))), tf.int32)
        pooling_region = inputs[:, start_h:end_h, start_w:end_w, :]
        pool_result = pooling_op(pooling_region, axis=(1, 2))
        result.append(pool_result)
    return result
  
  
  def spatial_pyramid_pool(inputs, dimensions=[2, 1], mode='max', implementation='kaiming'):
    """Performs spatial pyramid pooling (SPP) over the input. It will turn a 2D
    input of arbitrary size into an output of fixed dimenson. Hence, the
    convlutional part of a DNN can be connected to a dense part with a fixed
    number of nodes even if the dimensions of the input image are unknown.
  
    The pooling is performed over :math:`l` pooling levels.
    Each pooling level :math:`i` will create :math:`M_i` output features.
    :math:`M_i` is given by :math:`n_i * n_i`, with :math:`n_i` as the number
    of pooling operations per dimension level :math:`i`.
  
    The length of the parameter dimensions is the level of the spatial pyramid.
  
    Args:
        inputs: A 4D Tensor (B, H, W, C).
        dimensions: The list of :math:`n_i`'s that define the output dimension
        of each pooling level :math:`i`. The length of dimensions is the level of
        the spatial pyramid.
        mode: Pooling mode 'max' or 'avg'.
        implementation: The implementation to use, either 'kaiming' or 'fast'.
        kamming is the original implementation from the paper, and supports variable
        sizes of input vectors, which fast does not support.
  
    Returns:
        A fixed length vector representing the inputs.
  
    Notes:
        SPP should be inserted between the convolutional part of a DNN and it's
        dense part. Convolutions can be used for arbitrary input dimensions, but
        the size of their output will depend on their input dimensions.
        Connecting the output of the convolutional to the dense part then
        usually demands us to fix the dimensons of the network's input.
        The spatial pyramid pooling layer, however, allows us to leave
        the network input dimensions arbitrary.
        The advantage over a global pooling layer is the added robustness
        against object deformations due to the pooling on different scales.
    """
    pool_list = []
    if implementation == 'kaiming':
      for pool_dim in dimensions:
        pool_list += max_pool_2d_nxn_regions(inputs, pool_dim, mode)
    else:
      shape = inputs.get_shape().as_list()
      for d in dimensions:
        h = shape[1]
        w = shape[2]
        ph = np.ceil(h * 1.0 / d).astype(np.int32)
        pw = np.ceil(w * 1.0 / d).astype(np.int32)
        sh = np.floor(h * 1.0 / d + 1).astype(np.int32)
        sw = np.floor(w * 1.0 / d + 1).astype(np.int32)
        pool_result = tf.nn.max_pool(
            inputs, ksize=[1, ph, pw, 1], strides=[1, sh, sw, 1], padding='SAME')
        pool_list.append(tf.reshape(pool_result, [tf.shape(inputs)[0], -1]))
    return tf.concat(1, pool_list)
  
  
  def group_norm(x, filters=None, num_groups=8, epsilon=1e-5):
    x_shape = util.shape_list(x)
    if filters is None:
      filters = x_shape[-1]
    assert len(x_shape) == 4
    assert filters % num_groups == 0
    # Prepare variables.
    scale = tf.get_variable("group_norm_scale", [filters], initializer=tf.ones_initializer())
    bias = tf.get_variable("group_norm_bias", [filters], initializer=tf.zeros_initializer())
    epsilon, scale, bias = [tf.cast(t, x.dtype) for t in [epsilon, scale, bias]]
    # Reshape and compute group norm.
    x = tf.reshape(x, x_shape[:-1] + [num_groups, filters // num_groups])
    # Calculate mean and variance on heights, width, channels (not groups).
    mean, variance = tf.nn.moments(x, [1, 2, 4], keep_dims=True)
    norm_x = (x - mean) * tf.rsqrt(variance + epsilon)
    return tf.reshape(norm_x, x_shape) * scale + bias
  