                                  kernel_size=kernel_size,
                                  strides=strides,
                                  padding="same",
                                  name="{}_seperableconv2d".format(name),
                                  **kwargs)(inp)
          var_x = Activation("relu", name="{}_relu".format(name))(var_x)
          return var_x
  
  # <<< GAN V2.2 Blocks >>> #
  # TODO Merge these into NNBLock class when porting GAN2.2
 
</s> remove 
# <<< GAN V2.2 Blocks >>> #
# TODO Merge these into NNBLock class when porting GAN2.2


# Gan Constansts:
GAN22_CONV_INIT = "he_normal"
GAN22_REGULARIZER = 1e-4


# Gan Blocks:
def normalization(inp, norm="none", group="16"):
    """ GAN Normalization """
    if norm == "layernorm":
        var_x = GroupNormalization(group=group)(inp)
    elif norm == "batchnorm":
        var_x = BatchNormalization()(inp)
    elif norm == "groupnorm":
        var_x = GroupNormalization(group=16)(inp)
    elif norm == "instancenorm":
        var_x = InstanceNormalization()(inp)
    elif norm == "hybrid":
        if group % 2 == 1:
            raise ValueError("Output channels must be an even number for hybrid norm, "
                             "received {}.".format(group))
        filt = group
        var_x_0 = Lambda(lambda var_x: var_x[..., :filt // 2])(var_x)
        var_x_1 = Lambda(lambda var_x: var_x[..., filt // 2:])(var_x)
        var_x_0 = Conv2D(filt // 2,
                         kernel_size=1,
                         kernel_regularizer=regularizers.l2(GAN22_REGULARIZER),
                         kernel_initializer=GAN22_CONV_INIT)(var_x_0)
        var_x_1 = InstanceNormalization()(var_x_1)
        var_x = concatenate([var_x_0, var_x_1], axis=-1)
    else:
        var_x = inp
    return var_x


def upscale_ps(inp, filters, initializer, use_norm=False, norm="none"):
    """ GAN Upscaler - Pixel Shuffler """
    var_x = Conv2D(filters * 4,
                   kernel_size=3,
                   kernel_regularizer=regularizers.l2(GAN22_REGULARIZER),
                   kernel_initializer=initializer,
                   padding="same")(inp)
    var_x = LeakyReLU(0.2)(var_x)
    var_x = normalization(var_x, norm, filters) if use_norm else var_x
    var_x = PixelShuffler()(var_x)
    return var_x


def upscale_nn(inp, filters, use_norm=False, norm="none"):
    """ GAN Neural Network """
    var_x = UpSampling2D()(inp)
    var_x = reflect_padding_2d(var_x, 1)
    var_x = Conv2D(filters,
                   kernel_size=3,
                   kernel_regularizer=regularizers.l2(GAN22_REGULARIZER),
                   kernel_initializer="he_normal")(var_x)
    var_x = normalization(var_x, norm, filters) if use_norm else var_x
    return var_x


def reflect_padding_2d(inp, pad=1):
    """ GAN Reflect Padding (2D) """
    var_x = Lambda(lambda var_x: tf.pad(var_x,
                                        [[0, 0], [pad, pad], [pad, pad], [0, 0]],
                                        mode="REFLECT"))(inp)
    return var_x


def conv_gan(inp, filters, use_norm=False, strides=2, norm="none"):
    """ GAN Conv Block """
    var_x = Conv2D(filters,
                   kernel_size=3,
                   strides=strides,
                   kernel_regularizer=regularizers.l2(GAN22_REGULARIZER),
                   kernel_initializer=GAN22_CONV_INIT,
                   use_bias=False,
                   padding="same")(inp)
    var_x = Activation("relu")(var_x)
    var_x = normalization(var_x, norm, filters) if use_norm else var_x
    return var_x


def conv_d_gan(inp, filters, use_norm=False, norm="none"):
    """ GAN Discriminator Conv Block """
    var_x = inp
    var_x = Conv2D(filters,
                   kernel_size=4,
                   strides=2,
                   kernel_regularizer=regularizers.l2(GAN22_REGULARIZER),
                   kernel_initializer=GAN22_CONV_INIT,
                   use_bias=False,
                   padding="same")(var_x)
    var_x = LeakyReLU(alpha=0.2)(var_x)
    var_x = normalization(var_x, norm, filters) if use_norm else var_x
    return var_x


def res_block_gan(inp, filters, use_norm=False, norm="none"):
    """ GAN Res Block """
    var_x = Conv2D(filters,
                   kernel_size=3,
                   kernel_regularizer=regularizers.l2(GAN22_REGULARIZER),
                   kernel_initializer=GAN22_CONV_INIT,
                   use_bias=False,
                   padding="same")(inp)
    var_x = LeakyReLU(alpha=0.2)(var_x)
    var_x = normalization(var_x, norm, filters) if use_norm else var_x
    var_x = Conv2D(filters,
                   kernel_size=3,
                   kernel_regularizer=regularizers.l2(GAN22_REGULARIZER),
                   kernel_initializer=GAN22_CONV_INIT,
                   use_bias=False,
                   padding="same")(var_x)
    var_x = add([var_x, inp])
    var_x = LeakyReLU(alpha=0.2)(var_x)
    var_x = normalization(var_x, norm, filters) if use_norm else var_x
    return var_x


def self_attn_block(inp, n_c, squeeze_factor=8):
    """ GAN Self Attention Block
    Code borrows from https://github.com/taki0112/Self-Attention-GAN-Tensorflow
    """
    msg = "Input channels must be >= {}, recieved nc={}".format(squeeze_factor, n_c)
    assert n_c // squeeze_factor > 0, msg
    var_x = inp
    shape_x = var_x.get_shape().as_list()

    var_f = Conv2D(n_c // squeeze_factor, 1,
                   kernel_regularizer=regularizers.l2(GAN22_REGULARIZER))(var_x)
    var_g = Conv2D(n_c // squeeze_factor, 1,
                   kernel_regularizer=regularizers.l2(GAN22_REGULARIZER))(var_x)
    var_h = Conv2D(n_c, 1, kernel_regularizer=regularizers.l2(GAN22_REGULARIZER))(var_x)

    shape_f = var_f.get_shape().as_list()
    shape_g = var_g.get_shape().as_list()
    shape_h = var_h.get_shape().as_list()
    flat_f = Reshape((-1, shape_f[-1]))(var_f)
    flat_g = Reshape((-1, shape_g[-1]))(var_g)
    flat_h = Reshape((-1, shape_h[-1]))(var_h)

    var_s = Lambda(lambda var_x: K.batch_dot(var_x[0],
                                             Permute((2, 1))(var_x[1])))([flat_g, flat_f])

    beta = Softmax(axis=-1)(var_s)
    var_o = Lambda(lambda var_x: K.batch_dot(var_x[0], var_x[1]))([beta, flat_h])
    var_o = Reshape(shape_x[1:])(var_o)
    var_o = Scale()(var_o)

    out = add([var_o, inp])
    return out
 </s> add  </s> remove                        **kwargs)(inp)
 </s> add                        **kwargs)(input_tensor) </s> remove     def conv(self, inp, filters, kernel_size=5, strides=2, padding="same",
 </s> add     def conv(self, input_tensor, filters, kernel_size=5, strides=2, padding="same", </s> remove             self.switch_kernel_initializer(kwargs, original_init)
        var_x = Add()([var_x, inp])
 </s> add             self._switch_kernel_initializer(kwargs, original_init)
        var_x = Add()([var_x, input_tensor]) </s> remove     def conv_sep(self, inp, filters, kernel_size=5, strides=2, **kwargs):
        """ Seperable Convolution Layer """
        logger.debug("inp: %s, filters: %s, kernel_size: %s, strides: %s, kwargs: %s)",
                     inp, filters, kernel_size, strides, kwargs)
        name = self.get_name("separableconv2d_{}".format(inp.shape[1]))
        kwargs = self.set_default_initializer(kwargs)
 </s> add     def conv_sep(self, input_tensor, filters, kernel_size=5, strides=2, **kwargs):
        """ Seperable Convolution Layer.

        Parameters
        ----------
        input_tensor: tensor
            The input tensor to the layer
        filters: int
            The dimensionality of the output space (i.e. the number of output filters in the
            convolution)
        kernel_size: int, optional
            An integer or tuple/list of 2 integers, specifying the height and width of the 2D
            convolution window. Can be a single integer to specify the same value for all spatial
            dimensions. Default: 5
        strides: tuple or int, optional
            An integer or tuple/list of 2 integers, specifying the strides of the convolution along
            the height and width. Can be a single integer to specify the same value for all spatial
            dimensions. Default: `2`
        kwargs: dict
            Any additional Keras standard layer keyword arguments

        Returns
        -------
        tensor
            The output tensor from the Upscale layer
        """
        logger.debug("input_tensor: %s, filters: %s, kernel_size: %s, strides: %s, kwargs: %s)",
                     input_tensor, filters, kernel_size, strides, kwargs)
        name = self._get_name("separableconv2d_{}".format(input_tensor.shape[1]))
        kwargs = self._set_default_initializer(kwargs)
                                **kwargs)(input_tensor)
                                  kernel_size=kernel_size,
                                  strides=strides,
                                  padding="same",
                                  name="{}_seperableconv2d".format(name),
                                **kwargs)(input_tensor)
          var_x = Activation("relu", name="{}_relu".format(name))(var_x)
          return var_x
  
  # <<< GAN V2.2 Blocks >>> #
  # TODO Merge these into NNBLock class when porting GAN2.2
keep keep keep keep replace keep keep keep keep keep
 Simple backend unit tests (#1020)

* Add simple backend tests for lib.model
* Document lib.model
* Fix GMSD Loss for AMD
* Remove obsolete code from lib.model 
lib/model/nn_blocks.py