import logging

import tensorflow as tf
import tensorflow.keras.layers as layers
from tensorflow.keras.utils import get_custom_objects
from tensorflow.python.training.tracking import data_structures
from tensorflow.python.keras.layers.recurrent import DropoutRNNCellMixin
from tensorflow.python.keras.utils import tf_utils
from tensorflow.keras import backend as K
from tensorflow.python.util import nest
from tensorflow.python.ops import array_ops
from tensorflow.python.framework import tensor_shape
import numpy as np


################################################################################
# I STOLE THIS FROM THE TF-ADDONS REPO BECAUSE THEY DON'T PROVIDE A MAC VERSION
################################################################################
def register_keras_custom_object(cls):
    get_custom_objects()[cls.__name__] = cls
    return cls


@register_keras_custom_object
class GroupNormalization(tf.keras.layers.Layer):
    """Group normalization layer.

    Group Normalization divides the channels into groups and computes
    within each group the mean and variance for normalization.
    Empirically, its accuracy is more stable than batch norm in a wide
    range of small batch sizes, if learning rate is adjusted linearly
    with batch sizes.

    Relation to Layer Normalization:
    If the number of groups is set to 1, then this operation becomes identical
    to Layer Normalization.

    Relation to Instance Normalization:
    If the number of groups is set to the
    input dimension (number of groups is equal
    to number of channels), then this operation becomes
    identical to Instance Normalization.

    Arguments
        groups: Integer, the number of groups for Group Normalization.
            Can be in the range [1, N] where N is the input dimension.
            The input dimension must be divisible by the number of groups.
        axis: Integer, the axis that should be normalized.
        epsilon: Small float added to variance to avoid dividing by zero.
        center: If True, add offset of `beta` to normalized tensor.
            If False, `beta` is ignored.
        scale: If True, multiply by `gamma`.
            If False, `gamma` is not used.
        beta_initializer: Initializer for the beta weight.
        gamma_initializer: Initializer for the gamma weight.
        beta_regularizer: Optional regularizer for the beta weight.
        gamma_regularizer: Optional regularizer for the gamma weight.
        beta_constraint: Optional constraint for the beta weight.
        gamma_constraint: Optional constraint for the gamma weight.

    Input shape
        Arbitrary. Use the keyword argument `input_shape`
        (tuple of integers, does not include the samples axis)
        when using this layer as the first layer in a model.

    Output shape
        Same shape as input.
    References
        - [Group Normalization](https://arxiv.org/abs/1803.08494)
    """

    def __init__(self,
                 groups=2,
                 axis=-1,
                 epsilon=1e-5,
                 center=True,
                 scale=True,
                 beta_initializer='zeros',
                 gamma_initializer='ones',
                 beta_regularizer=None,
                 gamma_regularizer=None,
                 beta_constraint=None,
                 gamma_constraint=None,
                 **kwargs):
        super(GroupNormalization, self).__init__(**kwargs)
        self.supports_masking = True
        self.groups = groups
        self.axis = axis
        self.epsilon = epsilon
        self.center = center
        self.scale = scale
        self.beta_initializer = tf.keras.initializers.get(beta_initializer)
        self.gamma_initializer = tf.keras.initializers.get(gamma_initializer)
        self.beta_regularizer = tf.keras.regularizers.get(beta_regularizer)
        self.gamma_regularizer = tf.keras.regularizers.get(gamma_regularizer)
        self.beta_constraint = tf.keras.constraints.get(beta_constraint)
        self.gamma_constraint = tf.keras.constraints.get(gamma_constraint)
        self._check_axis()

    def build(self, input_shape):

        self._check_if_input_shape_is_none(input_shape)
        self._set_number_of_groups_for_instance_norm(input_shape)
        self._check_size_of_dimensions(input_shape)
        self._create_input_spec(input_shape)

        self._add_gamma_weight(input_shape)
        self._add_beta_weight(input_shape)
        self.built = True
        super(GroupNormalization, self).build(input_shape)

    def call(self, inputs):

        input_shape = tf.keras.backend.int_shape(inputs)
        tensor_input_shape = tf.shape(inputs)

        reshaped_inputs, group_shape = self._reshape_into_groups(
            inputs, input_shape, tensor_input_shape)

        normalized_inputs = self._apply_normalization(reshaped_inputs,
                                                      input_shape)

        outputs = tf.reshape(normalized_inputs, tensor_input_shape)

        return outputs

    def get_config(self):
        config = {
            'groups':
                self.groups,
            'axis':
                self.axis,
            'epsilon':
                self.epsilon,
            'center':
                self.center,
            'scale':
                self.scale,
            'beta_initializer':
                tf.keras.initializers.serialize(self.beta_initializer),
            'gamma_initializer':
                tf.keras.initializers.serialize(self.gamma_initializer),
            'beta_regularizer':
                tf.keras.regularizers.serialize(self.beta_regularizer),
            'gamma_regularizer':
                tf.keras.regularizers.serialize(self.gamma_regularizer),
            'beta_constraint':
                tf.keras.constraints.serialize(self.beta_constraint),
            'gamma_constraint':
                tf.keras.constraints.serialize(self.gamma_constraint)
        }
        base_config = super(GroupNormalization, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

    def compute_output_shape(self, input_shape):
        return input_shape

    def _reshape_into_groups(self, inputs, input_shape, tensor_input_shape):

        group_shape = [tensor_input_shape[i] for i in range(len(input_shape))]
        group_shape[self.axis] = input_shape[self.axis] // self.groups
        group_shape.insert(1, self.groups)
        group_shape = tf.stack(group_shape)
        reshaped_inputs = tf.reshape(inputs, group_shape)
        return reshaped_inputs, group_shape

    def _apply_normalization(self, reshaped_inputs, input_shape):

        group_shape = tf.keras.backend.int_shape(reshaped_inputs)
        group_reduction_axes = list(range(len(group_shape)))
        # Remember the ordering of the tensor is [batch, group , steps]. Jump
        # the first 2 to calculate the variance and the mean
        mean, variance = tf.nn.moments(
            reshaped_inputs, group_reduction_axes[2:], keepdims=True)

        gamma, beta = self._get_reshaped_weights(input_shape)
        normalized_inputs = tf.nn.batch_normalization(
            reshaped_inputs,
            mean=mean,
            variance=variance,
            scale=gamma,
            offset=beta,
            variance_epsilon=self.epsilon)
        return normalized_inputs

    def _get_reshaped_weights(self, input_shape):
        broadcast_shape = self._create_broadcast_shape(input_shape)
        gamma = None
        beta = None
        if self.scale:
            gamma = tf.reshape(self.gamma, broadcast_shape)

        if self.center:
            beta = tf.reshape(self.beta, broadcast_shape)
        return gamma, beta

    def _check_if_input_shape_is_none(self, input_shape):
        dim = input_shape[self.axis]
        if dim is None:
            raise ValueError('Axis ' + str(self.axis) + ' of '
                                                        'input tensor should have a defined dimension '
                                                        'but the layer received an input with shape ' +
                             str(input_shape) + '.')

    def _set_number_of_groups_for_instance_norm(self, input_shape):
        dim = input_shape[self.axis]

        if self.groups == -1:
            self.groups = dim

    def _check_size_of_dimensions(self, input_shape):

        dim = input_shape[self.axis]
        if dim < self.groups:
            raise ValueError(
                'Number of groups (' + str(self.groups) + ') cannot be '
                                                          'more than the number of channels (' + str(
                    dim) + ').')

        if dim % self.groups != 0:
            raise ValueError(
                'Number of groups (' + str(self.groups) + ') must be a '
                                                          'multiple of the number of channels (' + str(
                    dim) + ').')

    def _check_axis(self):

        if self.axis == 0:
            raise ValueError(
                "You are trying to normalize your batch axis. Do you want to "
                "use tf.layer.batch_normalization instead")

    def _create_input_spec(self, input_shape):

        dim = input_shape[self.axis]
        self.input_spec = tf.keras.layers.InputSpec(
            ndim=len(input_shape), axes={self.axis: dim})

    def _add_gamma_weight(self, input_shape):

        dim = input_shape[self.axis]
        shape = (dim,)

        if self.scale:
            self.gamma = self.add_weight(
                shape=shape,
                name='gamma',
                initializer=self.gamma_initializer,
                regularizer=self.gamma_regularizer,
                constraint=self.gamma_constraint)
        else:
            self.gamma = None

    def _add_beta_weight(self, input_shape):

        dim = input_shape[self.axis]
        shape = (dim,)

        if self.center:
            self.beta = self.add_weight(
                shape=shape,
                name='beta',
                initializer=self.beta_initializer,
                regularizer=self.beta_regularizer,
                constraint=self.beta_constraint)
        else:
            self.beta = None

    def _create_broadcast_shape(self, input_shape):
        broadcast_shape = [1] * len(input_shape)
        broadcast_shape[self.axis] = input_shape[self.axis] // self.groups
        broadcast_shape.insert(1, self.groups)
        return broadcast_shape


@register_keras_custom_object
class LayerNormalization(GroupNormalization):
    """Layer normalization layer.

    Layer Normalization is an specific case of ```GroupNormalization```since it
    normalizes all features of a layer. The Groupsize is 1.
    Empirically, its accuracy is more stable than batch norm in a wide
    range of small batch sizes, if learning rate is adjusted linearly
    with batch sizes.

    Arguments
        axis: Integer, the axis that should be normalized.
        epsilon: Small float added to variance to avoid dividing by zero.
        center: If True, add offset of `beta` to normalized tensor.
            If False, `beta` is ignored.
        scale: If True, multiply by `gamma`.
            If False, `gamma` is not used.
        beta_initializer: Initializer for the beta weight.
        gamma_initializer: Initializer for the gamma weight.
        beta_regularizer: Optional regularizer for the beta weight.
        gamma_regularizer: Optional regularizer for the gamma weight.
        beta_constraint: Optional constraint for the beta weight.
        gamma_constraint: Optional constraint for the gamma weight.

    Input shape
        Arbitrary. Use the keyword argument `input_shape`
        (tuple of integers, does not include the samples axis)
        when using this layer as the first layer in a model.

    Output shape
        Same shape as input.

    References
        - [Layer Normalization](https://arxiv.org/abs/1607.06450)
    """

    def __init__(self, **kwargs):
        if "groups" in kwargs:
            logging.warning("The given value for groups will be overwritten.")
        kwargs["groups"] = 1
        super(LayerNormalization, self).__init__(**kwargs)


@register_keras_custom_object
class InstanceNormalization(GroupNormalization):
    """Instance normalization layer.

    Instance Normalization is an specific case of ```GroupNormalization```since
    it normalizes all features of one channel. The Groupsize is equal to the
    channel size. Empirically, its accuracy is more stable than batch norm in a
    wide range of small batch sizes, if learning rate is adjusted linearly
    with batch sizes.

    Arguments
        axis: Integer, the axis that should be normalized.
        epsilon: Small float added to variance to avoid dividing by zero.
        center: If True, add offset of `beta` to normalized tensor.
            If False, `beta` is ignored.
        scale: If True, multiply by `gamma`.
            If False, `gamma` is not used.
        beta_initializer: Initializer for the beta weight.
        gamma_initializer: Initializer for the gamma weight.
        beta_regularizer: Optional regularizer for the beta weight.
        gamma_regularizer: Optional regularizer for the gamma weight.
        beta_constraint: Optional constraint for the beta weight.
        gamma_constraint: Optional constraint for the gamma weight.

    Input shape
        Arbitrary. Use the keyword argument `input_shape`
        (tuple of integers, does not include the samples axis)
        when using this layer as the first layer in a model.

    Output shape
        Same shape as input.

    References
        - [Instance Normalization: The Missing Ingredient for Fast Stylization]
        (https://arxiv.org/abs/1607.08022)
    """

    def __init__(self, **kwargs):
        if "groups" in kwargs:
            logging.warning("The given value for groups will be overwritten.")

        kwargs["groups"] = -1
        super(InstanceNormalization, self).__init__(**kwargs)


################################################################################
# END OF THIEVERY
################################################################################


@register_keras_custom_object
class ResidualBlock(layers.Layer):
    def __init__(self, units, filter_size, activation):
        super(ResidualBlock, self).__init__()
        self.l1 = layers.Conv2D(units, filter_size, padding="same")
        self.b1 = layers.BatchNormalization()
        self.l2 = layers.Conv2D(units, filter_size, padding="same")
        self.b2 = layers.BatchNormalization()
        self.units = units
        self.activation = activation
        self.shortcut = None

    def build(self, input_shape):
        if input_shape[-1] != self.units:
            self.shortcut = layers.Conv2D(self.units, 1, padding="same")
            self.shortcut.build(input_shape)

    def call(self, inputs):
        interm = self.activation(self.b1(self.l1(inputs)))
        interm = self.b2(self.l2(interm))
        if self.shortcut:
            inputs = self.shortcut(inputs)
        return self.activation(inputs + interm)


@register_keras_custom_object
class PhasedLSTMCell(DropoutRNNCellMixin, layers.Layer):
    """Cell class for the LSTM layer.
    Arguments:
      units: Positive integer, dimensionality of the output space.
      activation: Activation function to use.
        Default: hyperbolic tangent (`tanh`).
        If you pass `None`, no activation is applied
        (ie. "linear" activation: `a(x) = x`).
      recurrent_activation: Activation function to use
        for the recurrent step.
        Default: hard sigmoid (`hard_sigmoid`).
        If you pass `None`, no activation is applied
        (ie. "linear" activation: `a(x) = x`).
      use_bias: Boolean, whether the layer uses a bias vector.
      kernel_initializer: Initializer for the `kernel` weights matrix,
        used for the linear transformation of the inputs.
      recurrent_initializer: Initializer for the `recurrent_kernel`
        weights matrix,
        used for the linear transformation of the recurrent state.
      bias_initializer: Initializer for the bias vector.
      unit_forget_bias: Boolean.
        If True, add 1 to the bias of the forget gate at initialization.
        Setting it to true will also force `bias_initializer="zeros"`.
        This is recommended in [Jozefowicz et
          al.](http://www.jmlr.org/proceedings/papers/v37/jozefowicz15.pdf)
      kernel_regularizer: Regularizer function applied to
        the `kernel` weights matrix.
      recurrent_regularizer: Regularizer function applied to
        the `recurrent_kernel` weights matrix.
      bias_regularizer: Regularizer function applied to the bias vector.
      kernel_constraint: Constraint function applied to
        the `kernel` weights matrix.
      recurrent_constraint: Constraint function applied to
        the `recurrent_kernel` weights matrix.
      bias_constraint: Constraint function applied to the bias vector.
      dropout: Float between 0 and 1.
        Fraction of the units to drop for
        the linear transformation of the inputs.
      recurrent_dropout: Float between 0 and 1.
        Fraction of the units to drop for
        the linear transformation of the recurrent state.
      implementation: Implementation mode, either 1 or 2.
        Mode 1 will structure its operations as a larger number of
        smaller dot products and additions, whereas mode 2 will
        batch them into fewer, larger operations. These modes will
        have different performance profiles on different hardware and
        for different applications.
    Call arguments:
      inputs: A 2D tensor.
      states: List of state tensors corresponding to the previous timestep.
      training: Python boolean indicating whether the layer should behave in
        training mode or in inference mode. Only relevant when `dropout` or
        `recurrent_dropout` is used.
    """

    def __init__(self,
                 units,
                 activation=tf.nn.tanh,
                 recurrent_activation=tf.nn.sigmoid,
                 use_bias=True,
                 kernel_initializer='glorot_uniform',
                 recurrent_initializer='orthogonal',
                 bias_initializer='zeros',
                 unit_forget_bias=True,
                 kernel_regularizer=None,
                 recurrent_regularizer=None,
                 bias_regularizer=None,
                 kernel_constraint=None,
                 recurrent_constraint=None,
                 bias_constraint=None,
                 dropout=0.,
                 recurrent_dropout=0.,
                 implementation=2,
                 **kwargs):
        super(LSTMCell, self).__init__(**kwargs)
        self.units = units
        self.activation = activation
        self.recurrent_activation = recurrent_activation
        self.use_bias = use_bias

        self.kernel_initializer = tf.keras.initializers.get(kernel_initializer)
        self.recurrent_initializer = tf.keras.initializers.get(
            recurrent_initializer)
        self.bias_initializer = tf.keras.initializers.get(bias_initializer)
        self.unit_forget_bias = unit_forget_bias

        self.kernel_regularizer = tf.keras.regularizers.get(kernel_regularizer)
        self.recurrent_regularizer = tf.keras.regularizers.get(
            recurrent_regularizer)
        self.bias_regularizer = tf.keras.regularizers.get(bias_regularizer)

        self.kernel_constraint = tf.keras.constraints.get(kernel_constraint)
        self.recurrent_constraint = tf.keras.constraints.get(
            recurrent_constraint)
        self.bias_constraint = tf.keras.constraints.get(bias_constraint)

        self.dropout = min(1., max(0., dropout))
        self.recurrent_dropout = min(1., max(0., recurrent_dropout))
        self.implementation = implementation
        # tuple(_ListWrapper) was silently dropping list content in at least 2.7.10,
        # and fixed after 2.7.16. Converting the state_size to wrapper around
        # NoDependency(), so that the base_layer.__setattr__ will not convert it to
        # ListWrapper. Down the stream, self.states will be a list since it is
        # generated from nest.map_structure with list, and tuple(list) will work
        # properly.
        self.state_size = data_structures.NoDependency([self.units, self.units])
        self.output_size = self.units

    @tf_utils.shape_type_conversion
    def build(self, input_shape):
        input_dim = input_shape[-1]
        self.kernel = self.add_weight(
            shape=(input_dim, self.units * 4),
            name='kernel',
            initializer=self.kernel_initializer,
            regularizer=self.kernel_regularizer,
            constraint=self.kernel_constraint)
        self.recurrent_kernel = self.add_weight(
            shape=(self.units, self.units * 4),
            name='recurrent_kernel',
            initializer=self.recurrent_initializer,
            regularizer=self.recurrent_regularizer,
            constraint=self.recurrent_constraint)

        def exp_random_init(shape, dtype):
            rand = tf.random.uniform(shape, minval=1., maxval=6., dtype=dtype)
            return tf.exp(rand)

        def another_init(shape, dtype):
            return tf.random.uniform(shape, minval=0., maxval=1., dtype=dtype)

        self.periods = self.add_weight(
            shape=[self.units],
            name="timegate_period",
            initializer=exp_random_init,
            constraint=tf.nn.relu)
        self.phases = self.add_weight(
            shape=[self.units],
            name="timegate_phase",
            initializer=another_init,
            constraint=None)
        self.r_open = tf.constant(
            0.05 * np.ones([self.units], dtype=np.float32))
        self.leak = 0.001

        if self.use_bias:
            if self.unit_forget_bias:

                def bias_initializer(_, *args, **kwargs):
                    return K.concatenate([
                        self.bias_initializer((self.units,), *args, **kwargs),
                        tf.keras.initializers.Ones()((self.units,), *args,
                                                     **kwargs),
                        self.bias_initializer((self.units * 2,), *args,
                                              **kwargs),
                    ])
            else:
                bias_initializer = self.bias_initializer
            self.bias = self.add_weight(
                shape=(self.units * 4,),
                name='bias',
                initializer=bias_initializer,
                regularizer=self.bias_regularizer,
                constraint=self.bias_constraint)
        else:
            self.bias = None
        self.built = True

    def _compute_carry_and_output(self, x, h_tm1, c_tm1):
        """Computes carry and output using split kernels."""
        x_i, x_f, x_c, x_o = x
        h_tm1_i, h_tm1_f, h_tm1_c, h_tm1_o = h_tm1
        i = self.recurrent_activation(
            x_i + K.dot(h_tm1_i, self.recurrent_kernel[:, :self.units]))
        f = self.recurrent_activation(x_f + K.dot(
            h_tm1_f, self.recurrent_kernel[:, self.units:self.units * 2]))
        c = f * c_tm1 + i * self.activation(x_c + K.dot(
            h_tm1_c, self.recurrent_kernel[:, self.units * 2:self.units * 3]))
        o = self.recurrent_activation(
            x_o + K.dot(h_tm1_o, self.recurrent_kernel[:, self.units * 3:]))
        return c, o

    def _compute_carry_and_output_fused(self, z, c_tm1):
        """Computes carry and output using fused kernels."""
        z0, z1, z2, z3 = z
        i = self.recurrent_activation(z0)
        f = self.recurrent_activation(z1)
        c = f * c_tm1 + i * self.activation(z2)
        o = self.recurrent_activation(z3)
        return c, o

    def call(self, inputs, states, training=None):
        inputs, inp_time = inputs

        h_tm1 = states[0]  # previous memory state
        c_tm1 = states[1]  # previous carry state

        dp_mask = self.get_dropout_mask_for_cell(inputs, training, count=4)
        rec_dp_mask = self.get_recurrent_dropout_mask_for_cell(
            h_tm1, training, count=4)

        if self.implementation == 1:
            if 0 < self.dropout < 1.:
                inputs_i = inputs * dp_mask[0]
                inputs_f = inputs * dp_mask[1]
                inputs_c = inputs * dp_mask[2]
                inputs_o = inputs * dp_mask[3]
            else:
                inputs_i = inputs
                inputs_f = inputs
                inputs_c = inputs
                inputs_o = inputs
            k_i, k_f, k_c, k_o = array_ops.split(
                self.kernel, num_or_size_splits=4, axis=1)
            x_i = K.dot(inputs_i, k_i)
            x_f = K.dot(inputs_f, k_f)
            x_c = K.dot(inputs_c, k_c)
            x_o = K.dot(inputs_o, k_o)
            if self.use_bias:
                b_i, b_f, b_c, b_o = array_ops.split(
                    self.bias, num_or_size_splits=4, axis=0)
                x_i = K.bias_add(x_i, b_i)
                x_f = K.bias_add(x_f, b_f)
                x_c = K.bias_add(x_c, b_c)
                x_o = K.bias_add(x_o, b_o)

            if 0 < self.recurrent_dropout < 1.:
                h_tm1_i = h_tm1 * rec_dp_mask[0]
                h_tm1_f = h_tm1 * rec_dp_mask[1]
                h_tm1_c = h_tm1 * rec_dp_mask[2]
                h_tm1_o = h_tm1 * rec_dp_mask[3]
            else:
                h_tm1_i = h_tm1
                h_tm1_f = h_tm1
                h_tm1_c = h_tm1
                h_tm1_o = h_tm1
            x = (x_i, x_f, x_c, x_o)
            h_tm1 = (h_tm1_i, h_tm1_f, h_tm1_c, h_tm1_o)
            c, o = self._compute_carry_and_output(x, h_tm1, c_tm1)
        else:
            if 0. < self.dropout < 1.:
                inputs *= dp_mask[0]
            z = K.dot(inputs, self.kernel)
            if 0. < self.recurrent_dropout < 1.:
                h_tm1 *= rec_dp_mask[0]
            z += K.dot(h_tm1, self.recurrent_kernel)
            if self.use_bias:
                z = K.bias_add(z, self.bias)

            z = array_ops.split(z, num_or_size_splits=4, axis=1)
            c, o = self._compute_carry_and_output_fused(z, c_tm1)

        h = o * self.activation(c)

        # time gate stufffffff
        # time_gate = 1.  # use this for debugging (always open)
        # inp_time = 0.
        # inp_time is just batch. so we reshape to batch x 1
        # so that then broadcasting can do the rest to make everything
        # batch x features
        aux = tf.math.mod(inp_time[:, tf.newaxis] - self.phases,
                          self.periods) / self.periods

        p1 = tf.cast(tf.less(aux, 0.5 * self.r_open),
                     tf.float32) * 2 * aux / self.r_open
        p2 = tf.cast(tf.logical_and(tf.greater_equal(aux, 0.5 * self.r_open),
                                    tf.less(aux, self.r_open)), tf.float32) * (
                         2 - 2 * aux / self.r_open)
        p3 = tf.cast(tf.greater_equal(aux, self.r_open),
                     tf.float32) * self.leak * aux
        time_gate = p1 + p2 + p3

        h = time_gate * h + (1 - time_gate) * h_tm1
        c = time_gate * c + (1 - time_gate) * c_tm1
        return h, [h, c]

    def get_config(self):
        config = {
            'units':
                self.units,
            'activation':
                tf.keras.activations.serialize(self.activation),
            'recurrent_activation':
                tf.keras.activations.serialize(self.recurrent_activation),
            'use_bias':
                self.use_bias,
            'kernel_initializer':
                tf.keras.initializers.serialize(self.kernel_initializer),
            'recurrent_initializer':
                tf.keras.initializers.serialize(self.recurrent_initializer),
            'bias_initializer':
                tf.keras.initializers.serialize(self.bias_initializer),
            'unit_forget_bias':
                self.unit_forget_bias,
            'kernel_regularizer':
                tf.keras.regularizers.serialize(self.kernel_regularizer),
            'recurrent_regularizer':
                tf.keras.regularizers.serialize(self.recurrent_regularizer),
            'bias_regularizer':
                tf.keras.regularizers.serialize(self.bias_regularizer),
            'kernel_constraint':
                tf.keras.constraints.serialize(self.kernel_constraint),
            'recurrent_constraint':
                tf.keras.constraints.serialize(self.recurrent_constraint),
            'bias_constraint':
                tf.keras.constraints.serialize(self.bias_constraint),
            'dropout':
                self.dropout,
            'recurrent_dropout':
                self.recurrent_dropout,
            'implementation':
                self.implementation
        }
        base_config = super(PhasedLSTMCell, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

    def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
        return list(_generate_zero_filled_state_for_cell(
            self, inputs, batch_size, dtype))


def _generate_zero_filled_state_for_cell(cell, inputs, batch_size, dtype):
    if inputs is not None:
        batch_size = array_ops.shape(inputs)[0]
        dtype = inputs.dtype
    return _generate_zero_filled_state(batch_size, cell.state_size, dtype)


def _generate_zero_filled_state(batch_size_tensor, state_size, dtype):
    """Generate a zero filled tensor with shape [batch_size, state_size]."""
    if batch_size_tensor is None or dtype is None:
        raise ValueError(
            'batch_size and dtype cannot be None while constructing initial state: '
            'batch_size={}, dtype={}'.format(batch_size_tensor, dtype))

    def create_zeros(unnested_state_size):
        flat_dims = tensor_shape.as_shape(unnested_state_size).as_list()
        init_state_size = [batch_size_tensor] + flat_dims
        return array_ops.zeros(init_state_size, dtype=dtype)

    if nest.is_sequence(state_size):
        return nest.map_structure(create_zeros, state_size)
    else:
        return create_zeros(state_size)

0db1758f2dfee5160ea7aedcaab736bffbe849f8
<mask token>


@register_keras_custom_object
class InstanceNormalization(GroupNormalization):
    <mask token>

    def __init__(self, **kwargs):
        if 'groups' in kwargs:
            logging.warning('The given value for groups will be overwritten.')
        kwargs['groups'] = -1
        super(InstanceNormalization, self).__init__(**kwargs)


@register_keras_custom_object
class ResidualBlock(layers.Layer):

    def __init__(self, units, filter_size, activation):
        super(ResidualBlock, self).__init__()
        self.l1 = layers.Conv2D(units, filter_size, padding='same')
        self.b1 = layers.BatchNormalization()
        self.l2 = layers.Conv2D(units, filter_size, padding='same')
        self.b2 = layers.BatchNormalization()
        self.units = units
        self.activation = activation
        self.shortcut = None

    def build(self, input_shape):
        if input_shape[-1] != self.units:
            self.shortcut = layers.Conv2D(self.units, 1, padding='same')
            self.shortcut.build(input_shape)

    def call(self, inputs):
        interm = self.activation(self.b1(self.l1(inputs)))
        interm = self.b2(self.l2(interm))
        if self.shortcut:
            inputs = self.shortcut(inputs)
        return self.activation(inputs + interm)


@register_keras_custom_object
class PhasedLSTMCell(DropoutRNNCellMixin, layers.Layer):
    """Cell class for the LSTM layer.
    Arguments:
      units: Positive integer, dimensionality of the output space.
      activation: Activation function to use.
        Default: hyperbolic tangent (`tanh`).
        If you pass `None`, no activation is applied
        (ie. "linear" activation: `a(x) = x`).
      recurrent_activation: Activation function to use
        for the recurrent step.
        Default: hard sigmoid (`hard_sigmoid`).
        If you pass `None`, no activation is applied
        (ie. "linear" activation: `a(x) = x`).
      use_bias: Boolean, whether the layer uses a bias vector.
      kernel_initializer: Initializer for the `kernel` weights matrix,
        used for the linear transformation of the inputs.
      recurrent_initializer: Initializer for the `recurrent_kernel`
        weights matrix,
        used for the linear transformation of the recurrent state.
      bias_initializer: Initializer for the bias vector.
      unit_forget_bias: Boolean.
        If True, add 1 to the bias of the forget gate at initialization.
        Setting it to true will also force `bias_initializer="zeros"`.
        This is recommended in [Jozefowicz et
          al.](http://www.jmlr.org/proceedings/papers/v37/jozefowicz15.pdf)
      kernel_regularizer: Regularizer function applied to
        the `kernel` weights matrix.
      recurrent_regularizer: Regularizer function applied to
        the `recurrent_kernel` weights matrix.
      bias_regularizer: Regularizer function applied to the bias vector.
      kernel_constraint: Constraint function applied to
        the `kernel` weights matrix.
      recurrent_constraint: Constraint function applied to
        the `recurrent_kernel` weights matrix.
      bias_constraint: Constraint function applied to the bias vector.
      dropout: Float between 0 and 1.
        Fraction of the units to drop for
        the linear transformation of the inputs.
      recurrent_dropout: Float between 0 and 1.
        Fraction of the units to drop for
        the linear transformation of the recurrent state.
      implementation: Implementation mode, either 1 or 2.
        Mode 1 will structure its operations as a larger number of
        smaller dot products and additions, whereas mode 2 will
        batch them into fewer, larger operations. These modes will
        have different performance profiles on different hardware and
        for different applications.
    Call arguments:
      inputs: A 2D tensor.
      states: List of state tensors corresponding to the previous timestep.
      training: Python boolean indicating whether the layer should behave in
        training mode or in inference mode. Only relevant when `dropout` or
        `recurrent_dropout` is used.
    """

    def __init__(self, units, activation=tf.nn.tanh, recurrent_activation=
        tf.nn.sigmoid, use_bias=True, kernel_initializer='glorot_uniform',
        recurrent_initializer='orthogonal', bias_initializer='zeros',
        unit_forget_bias=True, kernel_regularizer=None,
        recurrent_regularizer=None, bias_regularizer=None,
        kernel_constraint=None, recurrent_constraint=None, bias_constraint=
        None, dropout=0.0, recurrent_dropout=0.0, implementation=2, **kwargs):
        super(LSTMCell, self).__init__(**kwargs)
        self.units = units
        self.activation = activation
        self.recurrent_activation = recurrent_activation
        self.use_bias = use_bias
        self.kernel_initializer = tf.keras.initializers.get(kernel_initializer)
        self.recurrent_initializer = tf.keras.initializers.get(
            recurrent_initializer)
        self.bias_initializer = tf.keras.initializers.get(bias_initializer)
        self.unit_forget_bias = unit_forget_bias
        self.kernel_regularizer = tf.keras.regularizers.get(kernel_regularizer)
        self.recurrent_regularizer = tf.keras.regularizers.get(
            recurrent_regularizer)
        self.bias_regularizer = tf.keras.regularizers.get(bias_regularizer)
        self.kernel_constraint = tf.keras.constraints.get(kernel_constraint)
        self.recurrent_constraint = tf.keras.constraints.get(
            recurrent_constraint)
        self.bias_constraint = tf.keras.constraints.get(bias_constraint)
        self.dropout = min(1.0, max(0.0, dropout))
        self.recurrent_dropout = min(1.0, max(0.0, recurrent_dropout))
        self.implementation = implementation
        self.state_size = data_structures.NoDependency([self.units, self.units]
            )
        self.output_size = self.units

    @tf_utils.shape_type_conversion
    def build(self, input_shape):
        input_dim = input_shape[-1]
        self.kernel = self.add_weight(shape=(input_dim, self.units * 4),
            name='kernel', initializer=self.kernel_initializer, regularizer
            =self.kernel_regularizer, constraint=self.kernel_constraint)
        self.recurrent_kernel = self.add_weight(shape=(self.units, self.
            units * 4), name='recurrent_kernel', initializer=self.
            recurrent_initializer, regularizer=self.recurrent_regularizer,
            constraint=self.recurrent_constraint)

        def exp_random_init(shape, dtype):
            rand = tf.random.uniform(shape, minval=1.0, maxval=6.0, dtype=dtype
                )
            return tf.exp(rand)

        def another_init(shape, dtype):
            return tf.random.uniform(shape, minval=0.0, maxval=1.0, dtype=dtype
                )
        self.periods = self.add_weight(shape=[self.units], name=
            'timegate_period', initializer=exp_random_init, constraint=tf.
            nn.relu)
        self.phases = self.add_weight(shape=[self.units], name=
            'timegate_phase', initializer=another_init, constraint=None)
        self.r_open = tf.constant(0.05 * np.ones([self.units], dtype=np.
            float32))
        self.leak = 0.001
        if self.use_bias:
            if self.unit_forget_bias:

                def bias_initializer(_, *args, **kwargs):
                    return K.concatenate([self.bias_initializer((self.units
                        ,), *args, **kwargs), tf.keras.initializers.Ones()(
                        (self.units,), *args, **kwargs), self.
                        bias_initializer((self.units * 2,), *args, **kwargs)])
            else:
                bias_initializer = self.bias_initializer
            self.bias = self.add_weight(shape=(self.units * 4,), name=
                'bias', initializer=bias_initializer, regularizer=self.
                bias_regularizer, constraint=self.bias_constraint)
        else:
            self.bias = None
        self.built = True

    def _compute_carry_and_output(self, x, h_tm1, c_tm1):
        """Computes carry and output using split kernels."""
        x_i, x_f, x_c, x_o = x
        h_tm1_i, h_tm1_f, h_tm1_c, h_tm1_o = h_tm1
        i = self.recurrent_activation(x_i + K.dot(h_tm1_i, self.
            recurrent_kernel[:, :self.units]))
        f = self.recurrent_activation(x_f + K.dot(h_tm1_f, self.
            recurrent_kernel[:, self.units:self.units * 2]))
        c = f * c_tm1 + i * self.activation(x_c + K.dot(h_tm1_c, self.
            recurrent_kernel[:, self.units * 2:self.units * 3]))
        o = self.recurrent_activation(x_o + K.dot(h_tm1_o, self.
            recurrent_kernel[:, self.units * 3:]))
        return c, o

    def _compute_carry_and_output_fused(self, z, c_tm1):
        """Computes carry and output using fused kernels."""
        z0, z1, z2, z3 = z
        i = self.recurrent_activation(z0)
        f = self.recurrent_activation(z1)
        c = f * c_tm1 + i * self.activation(z2)
        o = self.recurrent_activation(z3)
        return c, o

    def call(self, inputs, states, training=None):
        inputs, inp_time = inputs
        h_tm1 = states[0]
        c_tm1 = states[1]
        dp_mask = self.get_dropout_mask_for_cell(inputs, training, count=4)
        rec_dp_mask = self.get_recurrent_dropout_mask_for_cell(h_tm1,
            training, count=4)
        if self.implementation == 1:
            if 0 < self.dropout < 1.0:
                inputs_i = inputs * dp_mask[0]
                inputs_f = inputs * dp_mask[1]
                inputs_c = inputs * dp_mask[2]
                inputs_o = inputs * dp_mask[3]
            else:
                inputs_i = inputs
                inputs_f = inputs
                inputs_c = inputs
                inputs_o = inputs
            k_i, k_f, k_c, k_o = array_ops.split(self.kernel,
                num_or_size_splits=4, axis=1)
            x_i = K.dot(inputs_i, k_i)
            x_f = K.dot(inputs_f, k_f)
            x_c = K.dot(inputs_c, k_c)
            x_o = K.dot(inputs_o, k_o)
            if self.use_bias:
                b_i, b_f, b_c, b_o = array_ops.split(self.bias,
                    num_or_size_splits=4, axis=0)
                x_i = K.bias_add(x_i, b_i)
                x_f = K.bias_add(x_f, b_f)
                x_c = K.bias_add(x_c, b_c)
                x_o = K.bias_add(x_o, b_o)
            if 0 < self.recurrent_dropout < 1.0:
                h_tm1_i = h_tm1 * rec_dp_mask[0]
                h_tm1_f = h_tm1 * rec_dp_mask[1]
                h_tm1_c = h_tm1 * rec_dp_mask[2]
                h_tm1_o = h_tm1 * rec_dp_mask[3]
            else:
                h_tm1_i = h_tm1
                h_tm1_f = h_tm1
                h_tm1_c = h_tm1
                h_tm1_o = h_tm1
            x = x_i, x_f, x_c, x_o
            h_tm1 = h_tm1_i, h_tm1_f, h_tm1_c, h_tm1_o
            c, o = self._compute_carry_and_output(x, h_tm1, c_tm1)
        else:
            if 0.0 < self.dropout < 1.0:
                inputs *= dp_mask[0]
            z = K.dot(inputs, self.kernel)
            if 0.0 < self.recurrent_dropout < 1.0:
                h_tm1 *= rec_dp_mask[0]
            z += K.dot(h_tm1, self.recurrent_kernel)
            if self.use_bias:
                z = K.bias_add(z, self.bias)
            z = array_ops.split(z, num_or_size_splits=4, axis=1)
            c, o = self._compute_carry_and_output_fused(z, c_tm1)
        h = o * self.activation(c)
        aux = tf.math.mod(inp_time[:, tf.newaxis] - self.phases, self.periods
            ) / self.periods
        p1 = tf.cast(tf.less(aux, 0.5 * self.r_open), tf.float32
            ) * 2 * aux / self.r_open
        p2 = tf.cast(tf.logical_and(tf.greater_equal(aux, 0.5 * self.r_open
            ), tf.less(aux, self.r_open)), tf.float32) * (2 - 2 * aux /
            self.r_open)
        p3 = tf.cast(tf.greater_equal(aux, self.r_open), tf.float32
            ) * self.leak * aux
        time_gate = p1 + p2 + p3
        h = time_gate * h + (1 - time_gate) * h_tm1
        c = time_gate * c + (1 - time_gate) * c_tm1
        return h, [h, c]

    def get_config(self):
        config = {'units': self.units, 'activation': tf.keras.activations.
            serialize(self.activation), 'recurrent_activation': tf.keras.
            activations.serialize(self.recurrent_activation), 'use_bias':
            self.use_bias, 'kernel_initializer': tf.keras.initializers.
            serialize(self.kernel_initializer), 'recurrent_initializer': tf
            .keras.initializers.serialize(self.recurrent_initializer),
            'bias_initializer': tf.keras.initializers.serialize(self.
            bias_initializer), 'unit_forget_bias': self.unit_forget_bias,
            'kernel_regularizer': tf.keras.regularizers.serialize(self.
            kernel_regularizer), 'recurrent_regularizer': tf.keras.
            regularizers.serialize(self.recurrent_regularizer),
            'bias_regularizer': tf.keras.regularizers.serialize(self.
            bias_regularizer), 'kernel_constraint': tf.keras.constraints.
            serialize(self.kernel_constraint), 'recurrent_constraint': tf.
            keras.constraints.serialize(self.recurrent_constraint),
            'bias_constraint': tf.keras.constraints.serialize(self.
            bias_constraint), 'dropout': self.dropout, 'recurrent_dropout':
            self.recurrent_dropout, 'implementation': self.implementation}
        base_config = super(PhasedLSTMCell, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

    def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
        return list(_generate_zero_filled_state_for_cell(self, inputs,
            batch_size, dtype))


<mask token>

<mask token>


@register_keras_custom_object
class GroupNormalization(tf.keras.layers.Layer):
    <mask token>
    <mask token>

    def build(self, input_shape):
        self._check_if_input_shape_is_none(input_shape)
        self._set_number_of_groups_for_instance_norm(input_shape)
        self._check_size_of_dimensions(input_shape)
        self._create_input_spec(input_shape)
        self._add_gamma_weight(input_shape)
        self._add_beta_weight(input_shape)
        self.built = True
        super(GroupNormalization, self).build(input_shape)

    def call(self, inputs):
        input_shape = tf.keras.backend.int_shape(inputs)
        tensor_input_shape = tf.shape(inputs)
        reshaped_inputs, group_shape = self._reshape_into_groups(inputs,
            input_shape, tensor_input_shape)
        normalized_inputs = self._apply_normalization(reshaped_inputs,
            input_shape)
        outputs = tf.reshape(normalized_inputs, tensor_input_shape)
        return outputs

    def get_config(self):
        config = {'groups': self.groups, 'axis': self.axis, 'epsilon': self
            .epsilon, 'center': self.center, 'scale': self.scale,
            'beta_initializer': tf.keras.initializers.serialize(self.
            beta_initializer), 'gamma_initializer': tf.keras.initializers.
            serialize(self.gamma_initializer), 'beta_regularizer': tf.keras
            .regularizers.serialize(self.beta_regularizer),
            'gamma_regularizer': tf.keras.regularizers.serialize(self.
            gamma_regularizer), 'beta_constraint': tf.keras.constraints.
            serialize(self.beta_constraint), 'gamma_constraint': tf.keras.
            constraints.serialize(self.gamma_constraint)}
        base_config = super(GroupNormalization, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))
    <mask token>

    def _reshape_into_groups(self, inputs, input_shape, tensor_input_shape):
        group_shape = [tensor_input_shape[i] for i in range(len(input_shape))]
        group_shape[self.axis] = input_shape[self.axis] // self.groups
        group_shape.insert(1, self.groups)
        group_shape = tf.stack(group_shape)
        reshaped_inputs = tf.reshape(inputs, group_shape)
        return reshaped_inputs, group_shape

    def _apply_normalization(self, reshaped_inputs, input_shape):
        group_shape = tf.keras.backend.int_shape(reshaped_inputs)
        group_reduction_axes = list(range(len(group_shape)))
        mean, variance = tf.nn.moments(reshaped_inputs,
            group_reduction_axes[2:], keepdims=True)
        gamma, beta = self._get_reshaped_weights(input_shape)
        normalized_inputs = tf.nn.batch_normalization(reshaped_inputs, mean
            =mean, variance=variance, scale=gamma, offset=beta,
            variance_epsilon=self.epsilon)
        return normalized_inputs

    def _get_reshaped_weights(self, input_shape):
        broadcast_shape = self._create_broadcast_shape(input_shape)
        gamma = None
        beta = None
        if self.scale:
            gamma = tf.reshape(self.gamma, broadcast_shape)
        if self.center:
            beta = tf.reshape(self.beta, broadcast_shape)
        return gamma, beta
    <mask token>
    <mask token>

    def _check_size_of_dimensions(self, input_shape):
        dim = input_shape[self.axis]
        if dim < self.groups:
            raise ValueError('Number of groups (' + str(self.groups) +
                ') cannot be more than the number of channels (' + str(dim) +
                ').')
        if dim % self.groups != 0:
            raise ValueError('Number of groups (' + str(self.groups) +
                ') must be a multiple of the number of channels (' + str(
                dim) + ').')
    <mask token>
    <mask token>
    <mask token>
    <mask token>
    <mask token>


@register_keras_custom_object
class LayerNormalization(GroupNormalization):
    """Layer normalization layer.

    Layer Normalization is an specific case of ```GroupNormalization```since it
    normalizes all features of a layer. The Groupsize is 1.
    Empirically, its accuracy is more stable than batch norm in a wide
    range of small batch sizes, if learning rate is adjusted linearly
    with batch sizes.

    Arguments
        axis: Integer, the axis that should be normalized.
        epsilon: Small float added to variance to avoid dividing by zero.
        center: If True, add offset of `beta` to normalized tensor.
            If False, `beta` is ignored.
        scale: If True, multiply by `gamma`.
            If False, `gamma` is not used.
        beta_initializer: Initializer for the beta weight.
        gamma_initializer: Initializer for the gamma weight.
        beta_regularizer: Optional regularizer for the beta weight.
        gamma_regularizer: Optional regularizer for the gamma weight.
        beta_constraint: Optional constraint for the beta weight.
        gamma_constraint: Optional constraint for the gamma weight.

    Input shape
        Arbitrary. Use the keyword argument `input_shape`
        (tuple of integers, does not include the samples axis)
        when using this layer as the first layer in a model.

    Output shape
        Same shape as input.

    References
        - [Layer Normalization](https://arxiv.org/abs/1607.06450)
    """

    def __init__(self, **kwargs):
        if 'groups' in kwargs:
            logging.warning('The given value for groups will be overwritten.')
        kwargs['groups'] = 1
        super(LayerNormalization, self).__init__(**kwargs)


@register_keras_custom_object
class InstanceNormalization(GroupNormalization):
    """Instance normalization layer.

    Instance Normalization is an specific case of ```GroupNormalization```since
    it normalizes all features of one channel. The Groupsize is equal to the
    channel size. Empirically, its accuracy is more stable than batch norm in a
    wide range of small batch sizes, if learning rate is adjusted linearly
    with batch sizes.

    Arguments
        axis: Integer, the axis that should be normalized.
        epsilon: Small float added to variance to avoid dividing by zero.
        center: If True, add offset of `beta` to normalized tensor.
            If False, `beta` is ignored.
        scale: If True, multiply by `gamma`.
            If False, `gamma` is not used.
        beta_initializer: Initializer for the beta weight.
        gamma_initializer: Initializer for the gamma weight.
        beta_regularizer: Optional regularizer for the beta weight.
        gamma_regularizer: Optional regularizer for the gamma weight.
        beta_constraint: Optional constraint for the beta weight.
        gamma_constraint: Optional constraint for the gamma weight.

    Input shape
        Arbitrary. Use the keyword argument `input_shape`
        (tuple of integers, does not include the samples axis)
        when using this layer as the first layer in a model.

    Output shape
        Same shape as input.

    References
        - [Instance Normalization: The Missing Ingredient for Fast Stylization]
        (https://arxiv.org/abs/1607.08022)
    """

    def __init__(self, **kwargs):
        if 'groups' in kwargs:
            logging.warning('The given value for groups will be overwritten.')
        kwargs['groups'] = -1
        super(InstanceNormalization, self).__init__(**kwargs)


@register_keras_custom_object
class ResidualBlock(layers.Layer):

    def __init__(self, units, filter_size, activation):
        super(ResidualBlock, self).__init__()
        self.l1 = layers.Conv2D(units, filter_size, padding='same')
        self.b1 = layers.BatchNormalization()
        self.l2 = layers.Conv2D(units, filter_size, padding='same')
        self.b2 = layers.BatchNormalization()
        self.units = units
        self.activation = activation
        self.shortcut = None

    def build(self, input_shape):
        if input_shape[-1] != self.units:
            self.shortcut = layers.Conv2D(self.units, 1, padding='same')
            self.shortcut.build(input_shape)

    def call(self, inputs):
        interm = self.activation(self.b1(self.l1(inputs)))
        interm = self.b2(self.l2(interm))
        if self.shortcut:
            inputs = self.shortcut(inputs)
        return self.activation(inputs + interm)


@register_keras_custom_object
class PhasedLSTMCell(DropoutRNNCellMixin, layers.Layer):
    """Cell class for the LSTM layer.
    Arguments:
      units: Positive integer, dimensionality of the output space.
      activation: Activation function to use.
        Default: hyperbolic tangent (`tanh`).
        If you pass `None`, no activation is applied
        (ie. "linear" activation: `a(x) = x`).
      recurrent_activation: Activation function to use
        for the recurrent step.
        Default: hard sigmoid (`hard_sigmoid`).
        If you pass `None`, no activation is applied
        (ie. "linear" activation: `a(x) = x`).
      use_bias: Boolean, whether the layer uses a bias vector.
      kernel_initializer: Initializer for the `kernel` weights matrix,
        used for the linear transformation of the inputs.
      recurrent_initializer: Initializer for the `recurrent_kernel`
        weights matrix,
        used for the linear transformation of the recurrent state.
      bias_initializer: Initializer for the bias vector.
      unit_forget_bias: Boolean.
        If True, add 1 to the bias of the forget gate at initialization.
        Setting it to true will also force `bias_initializer="zeros"`.
        This is recommended in [Jozefowicz et
          al.](http://www.jmlr.org/proceedings/papers/v37/jozefowicz15.pdf)
      kernel_regularizer: Regularizer function applied to
        the `kernel` weights matrix.
      recurrent_regularizer: Regularizer function applied to
        the `recurrent_kernel` weights matrix.
      bias_regularizer: Regularizer function applied to the bias vector.
      kernel_constraint: Constraint function applied to
        the `kernel` weights matrix.
      recurrent_constraint: Constraint function applied to
        the `recurrent_kernel` weights matrix.
      bias_constraint: Constraint function applied to the bias vector.
      dropout: Float between 0 and 1.
        Fraction of the units to drop for
        the linear transformation of the inputs.
      recurrent_dropout: Float between 0 and 1.
        Fraction of the units to drop for
        the linear transformation of the recurrent state.
      implementation: Implementation mode, either 1 or 2.
        Mode 1 will structure its operations as a larger number of
        smaller dot products and additions, whereas mode 2 will
        batch them into fewer, larger operations. These modes will
        have different performance profiles on different hardware and
        for different applications.
    Call arguments:
      inputs: A 2D tensor.
      states: List of state tensors corresponding to the previous timestep.
      training: Python boolean indicating whether the layer should behave in
        training mode or in inference mode. Only relevant when `dropout` or
        `recurrent_dropout` is used.
    """

    def __init__(self, units, activation=tf.nn.tanh, recurrent_activation=
        tf.nn.sigmoid, use_bias=True, kernel_initializer='glorot_uniform',
        recurrent_initializer='orthogonal', bias_initializer='zeros',
        unit_forget_bias=True, kernel_regularizer=None,
        recurrent_regularizer=None, bias_regularizer=None,
        kernel_constraint=None, recurrent_constraint=None, bias_constraint=
        None, dropout=0.0, recurrent_dropout=0.0, implementation=2, **kwargs):
        super(LSTMCell, self).__init__(**kwargs)
        self.units = units
        self.activation = activation
        self.recurrent_activation = recurrent_activation
        self.use_bias = use_bias
        self.kernel_initializer = tf.keras.initializers.get(kernel_initializer)
        self.recurrent_initializer = tf.keras.initializers.get(
            recurrent_initializer)
        self.bias_initializer = tf.keras.initializers.get(bias_initializer)
        self.unit_forget_bias = unit_forget_bias
        self.kernel_regularizer = tf.keras.regularizers.get(kernel_regularizer)
        self.recurrent_regularizer = tf.keras.regularizers.get(
            recurrent_regularizer)
        self.bias_regularizer = tf.keras.regularizers.get(bias_regularizer)
        self.kernel_constraint = tf.keras.constraints.get(kernel_constraint)
        self.recurrent_constraint = tf.keras.constraints.get(
            recurrent_constraint)
        self.bias_constraint = tf.keras.constraints.get(bias_constraint)
        self.dropout = min(1.0, max(0.0, dropout))
        self.recurrent_dropout = min(1.0, max(0.0, recurrent_dropout))
        self.implementation = implementation
        self.state_size = data_structures.NoDependency([self.units, self.units]
            )
        self.output_size = self.units

    @tf_utils.shape_type_conversion
    def build(self, input_shape):
        input_dim = input_shape[-1]
        self.kernel = self.add_weight(shape=(input_dim, self.units * 4),
            name='kernel', initializer=self.kernel_initializer, regularizer
            =self.kernel_regularizer, constraint=self.kernel_constraint)
        self.recurrent_kernel = self.add_weight(shape=(self.units, self.
            units * 4), name='recurrent_kernel', initializer=self.
            recurrent_initializer, regularizer=self.recurrent_regularizer,
            constraint=self.recurrent_constraint)

        def exp_random_init(shape, dtype):
            rand = tf.random.uniform(shape, minval=1.0, maxval=6.0, dtype=dtype
                )
            return tf.exp(rand)

        def another_init(shape, dtype):
            return tf.random.uniform(shape, minval=0.0, maxval=1.0, dtype=dtype
                )
        self.periods = self.add_weight(shape=[self.units], name=
            'timegate_period', initializer=exp_random_init, constraint=tf.
            nn.relu)
        self.phases = self.add_weight(shape=[self.units], name=
            'timegate_phase', initializer=another_init, constraint=None)
        self.r_open = tf.constant(0.05 * np.ones([self.units], dtype=np.
            float32))
        self.leak = 0.001
        if self.use_bias:
            if self.unit_forget_bias:

                def bias_initializer(_, *args, **kwargs):
                    return K.concatenate([self.bias_initializer((self.units
                        ,), *args, **kwargs), tf.keras.initializers.Ones()(
                        (self.units,), *args, **kwargs), self.
                        bias_initializer((self.units * 2,), *args, **kwargs)])
            else:
                bias_initializer = self.bias_initializer
            self.bias = self.add_weight(shape=(self.units * 4,), name=
                'bias', initializer=bias_initializer, regularizer=self.
                bias_regularizer, constraint=self.bias_constraint)
        else:
            self.bias = None
        self.built = True

    def _compute_carry_and_output(self, x, h_tm1, c_tm1):
        """Computes carry and output using split kernels."""
        x_i, x_f, x_c, x_o = x
        h_tm1_i, h_tm1_f, h_tm1_c, h_tm1_o = h_tm1
        i = self.recurrent_activation(x_i + K.dot(h_tm1_i, self.
            recurrent_kernel[:, :self.units]))
        f = self.recurrent_activation(x_f + K.dot(h_tm1_f, self.
            recurrent_kernel[:, self.units:self.units * 2]))
        c = f * c_tm1 + i * self.activation(x_c + K.dot(h_tm1_c, self.
            recurrent_kernel[:, self.units * 2:self.units * 3]))
        o = self.recurrent_activation(x_o + K.dot(h_tm1_o, self.
            recurrent_kernel[:, self.units * 3:]))
        return c, o

    def _compute_carry_and_output_fused(self, z, c_tm1):
        """Computes carry and output using fused kernels."""
        z0, z1, z2, z3 = z
        i = self.recurrent_activation(z0)
        f = self.recurrent_activation(z1)
        c = f * c_tm1 + i * self.activation(z2)
        o = self.recurrent_activation(z3)
        return c, o

    def call(self, inputs, states, training=None):
        inputs, inp_time = inputs
        h_tm1 = states[0]
        c_tm1 = states[1]
        dp_mask = self.get_dropout_mask_for_cell(inputs, training, count=4)
        rec_dp_mask = self.get_recurrent_dropout_mask_for_cell(h_tm1,
            training, count=4)
        if self.implementation == 1:
            if 0 < self.dropout < 1.0:
                inputs_i = inputs * dp_mask[0]
                inputs_f = inputs * dp_mask[1]
                inputs_c = inputs * dp_mask[2]
                inputs_o = inputs * dp_mask[3]
            else:
                inputs_i = inputs
                inputs_f = inputs
                inputs_c = inputs
                inputs_o = inputs
            k_i, k_f, k_c, k_o = array_ops.split(self.kernel,
                num_or_size_splits=4, axis=1)
            x_i = K.dot(inputs_i, k_i)
            x_f = K.dot(inputs_f, k_f)
            x_c = K.dot(inputs_c, k_c)
            x_o = K.dot(inputs_o, k_o)
            if self.use_bias:
                b_i, b_f, b_c, b_o = array_ops.split(self.bias,
                    num_or_size_splits=4, axis=0)
                x_i = K.bias_add(x_i, b_i)
                x_f = K.bias_add(x_f, b_f)
                x_c = K.bias_add(x_c, b_c)
                x_o = K.bias_add(x_o, b_o)
            if 0 < self.recurrent_dropout < 1.0:
                h_tm1_i = h_tm1 * rec_dp_mask[0]
                h_tm1_f = h_tm1 * rec_dp_mask[1]
                h_tm1_c = h_tm1 * rec_dp_mask[2]
                h_tm1_o = h_tm1 * rec_dp_mask[3]
            else:
                h_tm1_i = h_tm1
                h_tm1_f = h_tm1
                h_tm1_c = h_tm1
                h_tm1_o = h_tm1
            x = x_i, x_f, x_c, x_o
            h_tm1 = h_tm1_i, h_tm1_f, h_tm1_c, h_tm1_o
            c, o = self._compute_carry_and_output(x, h_tm1, c_tm1)
        else:
            if 0.0 < self.dropout < 1.0:
                inputs *= dp_mask[0]
            z = K.dot(inputs, self.kernel)
            if 0.0 < self.recurrent_dropout < 1.0:
                h_tm1 *= rec_dp_mask[0]
            z += K.dot(h_tm1, self.recurrent_kernel)
            if self.use_bias:
                z = K.bias_add(z, self.bias)
            z = array_ops.split(z, num_or_size_splits=4, axis=1)
            c, o = self._compute_carry_and_output_fused(z, c_tm1)
        h = o * self.activation(c)
        aux = tf.math.mod(inp_time[:, tf.newaxis] - self.phases, self.periods
            ) / self.periods
        p1 = tf.cast(tf.less(aux, 0.5 * self.r_open), tf.float32
            ) * 2 * aux / self.r_open
        p2 = tf.cast(tf.logical_and(tf.greater_equal(aux, 0.5 * self.r_open
            ), tf.less(aux, self.r_open)), tf.float32) * (2 - 2 * aux /
            self.r_open)
        p3 = tf.cast(tf.greater_equal(aux, self.r_open), tf.float32
            ) * self.leak * aux
        time_gate = p1 + p2 + p3
        h = time_gate * h + (1 - time_gate) * h_tm1
        c = time_gate * c + (1 - time_gate) * c_tm1
        return h, [h, c]

    def get_config(self):
        config = {'units': self.units, 'activation': tf.keras.activations.
            serialize(self.activation), 'recurrent_activation': tf.keras.
            activations.serialize(self.recurrent_activation), 'use_bias':
            self.use_bias, 'kernel_initializer': tf.keras.initializers.
            serialize(self.kernel_initializer), 'recurrent_initializer': tf
            .keras.initializers.serialize(self.recurrent_initializer),
            'bias_initializer': tf.keras.initializers.serialize(self.
            bias_initializer), 'unit_forget_bias': self.unit_forget_bias,
            'kernel_regularizer': tf.keras.regularizers.serialize(self.
            kernel_regularizer), 'recurrent_regularizer': tf.keras.
            regularizers.serialize(self.recurrent_regularizer),
            'bias_regularizer': tf.keras.regularizers.serialize(self.
            bias_regularizer), 'kernel_constraint': tf.keras.constraints.
            serialize(self.kernel_constraint), 'recurrent_constraint': tf.
            keras.constraints.serialize(self.recurrent_constraint),
            'bias_constraint': tf.keras.constraints.serialize(self.
            bias_constraint), 'dropout': self.dropout, 'recurrent_dropout':
            self.recurrent_dropout, 'implementation': self.implementation}
        base_config = super(PhasedLSTMCell, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

    def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
        return list(_generate_zero_filled_state_for_cell(self, inputs,
            batch_size, dtype))


<mask token>

<mask token>


@register_keras_custom_object
class GroupNormalization(tf.keras.layers.Layer):
    <mask token>
    <mask token>

    def build(self, input_shape):
        self._check_if_input_shape_is_none(input_shape)
        self._set_number_of_groups_for_instance_norm(input_shape)
        self._check_size_of_dimensions(input_shape)
        self._create_input_spec(input_shape)
        self._add_gamma_weight(input_shape)
        self._add_beta_weight(input_shape)
        self.built = True
        super(GroupNormalization, self).build(input_shape)

    def call(self, inputs):
        input_shape = tf.keras.backend.int_shape(inputs)
        tensor_input_shape = tf.shape(inputs)
        reshaped_inputs, group_shape = self._reshape_into_groups(inputs,
            input_shape, tensor_input_shape)
        normalized_inputs = self._apply_normalization(reshaped_inputs,
            input_shape)
        outputs = tf.reshape(normalized_inputs, tensor_input_shape)
        return outputs

    def get_config(self):
        config = {'groups': self.groups, 'axis': self.axis, 'epsilon': self
            .epsilon, 'center': self.center, 'scale': self.scale,
            'beta_initializer': tf.keras.initializers.serialize(self.
            beta_initializer), 'gamma_initializer': tf.keras.initializers.
            serialize(self.gamma_initializer), 'beta_regularizer': tf.keras
            .regularizers.serialize(self.beta_regularizer),
            'gamma_regularizer': tf.keras.regularizers.serialize(self.
            gamma_regularizer), 'beta_constraint': tf.keras.constraints.
            serialize(self.beta_constraint), 'gamma_constraint': tf.keras.
            constraints.serialize(self.gamma_constraint)}
        base_config = super(GroupNormalization, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))
    <mask token>

    def _reshape_into_groups(self, inputs, input_shape, tensor_input_shape):
        group_shape = [tensor_input_shape[i] for i in range(len(input_shape))]
        group_shape[self.axis] = input_shape[self.axis] // self.groups
        group_shape.insert(1, self.groups)
        group_shape = tf.stack(group_shape)
        reshaped_inputs = tf.reshape(inputs, group_shape)
        return reshaped_inputs, group_shape

    def _apply_normalization(self, reshaped_inputs, input_shape):
        group_shape = tf.keras.backend.int_shape(reshaped_inputs)
        group_reduction_axes = list(range(len(group_shape)))
        mean, variance = tf.nn.moments(reshaped_inputs,
            group_reduction_axes[2:], keepdims=True)
        gamma, beta = self._get_reshaped_weights(input_shape)
        normalized_inputs = tf.nn.batch_normalization(reshaped_inputs, mean
            =mean, variance=variance, scale=gamma, offset=beta,
            variance_epsilon=self.epsilon)
        return normalized_inputs

    def _get_reshaped_weights(self, input_shape):
        broadcast_shape = self._create_broadcast_shape(input_shape)
        gamma = None
        beta = None
        if self.scale:
            gamma = tf.reshape(self.gamma, broadcast_shape)
        if self.center:
            beta = tf.reshape(self.beta, broadcast_shape)
        return gamma, beta
    <mask token>
    <mask token>

    def _check_size_of_dimensions(self, input_shape):
        dim = input_shape[self.axis]
        if dim < self.groups:
            raise ValueError('Number of groups (' + str(self.groups) +
                ') cannot be more than the number of channels (' + str(dim) +
                ').')
        if dim % self.groups != 0:
            raise ValueError('Number of groups (' + str(self.groups) +
                ') must be a multiple of the number of channels (' + str(
                dim) + ').')

    def _check_axis(self):
        if self.axis == 0:
            raise ValueError(
                'You are trying to normalize your batch axis. Do you want to use tf.layer.batch_normalization instead'
                )
    <mask token>
    <mask token>
    <mask token>

    def _create_broadcast_shape(self, input_shape):
        broadcast_shape = [1] * len(input_shape)
        broadcast_shape[self.axis] = input_shape[self.axis] // self.groups
        broadcast_shape.insert(1, self.groups)
        return broadcast_shape


@register_keras_custom_object
class LayerNormalization(GroupNormalization):
    """Layer normalization layer.

    Layer Normalization is an specific case of ```GroupNormalization```since it
    normalizes all features of a layer. The Groupsize is 1.
    Empirically, its accuracy is more stable than batch norm in a wide
    range of small batch sizes, if learning rate is adjusted linearly
    with batch sizes.

    Arguments
        axis: Integer, the axis that should be normalized.
        epsilon: Small float added to variance to avoid dividing by zero.
        center: If True, add offset of `beta` to normalized tensor.
            If False, `beta` is ignored.
        scale: If True, multiply by `gamma`.
            If False, `gamma` is not used.
        beta_initializer: Initializer for the beta weight.
        gamma_initializer: Initializer for the gamma weight.
        beta_regularizer: Optional regularizer for the beta weight.
        gamma_regularizer: Optional regularizer for the gamma weight.
        beta_constraint: Optional constraint for the beta weight.
        gamma_constraint: Optional constraint for the gamma weight.

    Input shape
        Arbitrary. Use the keyword argument `input_shape`
        (tuple of integers, does not include the samples axis)
        when using this layer as the first layer in a model.

    Output shape
        Same shape as input.

    References
        - [Layer Normalization](https://arxiv.org/abs/1607.06450)
    """

    def __init__(self, **kwargs):
        if 'groups' in kwargs:
            logging.warning('The given value for groups will be overwritten.')
        kwargs['groups'] = 1
        super(LayerNormalization, self).__init__(**kwargs)


@register_keras_custom_object
class InstanceNormalization(GroupNormalization):
    """Instance normalization layer.

    Instance Normalization is an specific case of ```GroupNormalization```since
    it normalizes all features of one channel. The Groupsize is equal to the
    channel size. Empirically, its accuracy is more stable than batch norm in a
    wide range of small batch sizes, if learning rate is adjusted linearly
    with batch sizes.

    Arguments
        axis: Integer, the axis that should be normalized.
        epsilon: Small float added to variance to avoid dividing by zero.
        center: If True, add offset of `beta` to normalized tensor.
            If False, `beta` is ignored.
        scale: If True, multiply by `gamma`.
            If False, `gamma` is not used.
        beta_initializer: Initializer for the beta weight.
        gamma_initializer: Initializer for the gamma weight.
        beta_regularizer: Optional regularizer for the beta weight.
        gamma_regularizer: Optional regularizer for the gamma weight.
        beta_constraint: Optional constraint for the beta weight.
        gamma_constraint: Optional constraint for the gamma weight.

    Input shape
        Arbitrary. Use the keyword argument `input_shape`
        (tuple of integers, does not include the samples axis)
        when using this layer as the first layer in a model.

    Output shape
        Same shape as input.

    References
        - [Instance Normalization: The Missing Ingredient for Fast Stylization]
        (https://arxiv.org/abs/1607.08022)
    """

    def __init__(self, **kwargs):
        if 'groups' in kwargs:
            logging.warning('The given value for groups will be overwritten.')
        kwargs['groups'] = -1
        super(InstanceNormalization, self).__init__(**kwargs)


@register_keras_custom_object
class ResidualBlock(layers.Layer):

    def __init__(self, units, filter_size, activation):
        super(ResidualBlock, self).__init__()
        self.l1 = layers.Conv2D(units, filter_size, padding='same')
        self.b1 = layers.BatchNormalization()
        self.l2 = layers.Conv2D(units, filter_size, padding='same')
        self.b2 = layers.BatchNormalization()
        self.units = units
        self.activation = activation
        self.shortcut = None

    def build(self, input_shape):
        if input_shape[-1] != self.units:
            self.shortcut = layers.Conv2D(self.units, 1, padding='same')
            self.shortcut.build(input_shape)

    def call(self, inputs):
        interm = self.activation(self.b1(self.l1(inputs)))
        interm = self.b2(self.l2(interm))
        if self.shortcut:
            inputs = self.shortcut(inputs)
        return self.activation(inputs + interm)


@register_keras_custom_object
class PhasedLSTMCell(DropoutRNNCellMixin, layers.Layer):
    """Cell class for the LSTM layer.
    Arguments:
      units: Positive integer, dimensionality of the output space.
      activation: Activation function to use.
        Default: hyperbolic tangent (`tanh`).
        If you pass `None`, no activation is applied
        (ie. "linear" activation: `a(x) = x`).
      recurrent_activation: Activation function to use
        for the recurrent step.
        Default: hard sigmoid (`hard_sigmoid`).
        If you pass `None`, no activation is applied
        (ie. "linear" activation: `a(x) = x`).
      use_bias: Boolean, whether the layer uses a bias vector.
      kernel_initializer: Initializer for the `kernel` weights matrix,
        used for the linear transformation of the inputs.
      recurrent_initializer: Initializer for the `recurrent_kernel`
        weights matrix,
        used for the linear transformation of the recurrent state.
      bias_initializer: Initializer for the bias vector.
      unit_forget_bias: Boolean.
        If True, add 1 to the bias of the forget gate at initialization.
        Setting it to true will also force `bias_initializer="zeros"`.
        This is recommended in [Jozefowicz et
          al.](http://www.jmlr.org/proceedings/papers/v37/jozefowicz15.pdf)
      kernel_regularizer: Regularizer function applied to
        the `kernel` weights matrix.
      recurrent_regularizer: Regularizer function applied to
        the `recurrent_kernel` weights matrix.
      bias_regularizer: Regularizer function applied to the bias vector.
      kernel_constraint: Constraint function applied to
        the `kernel` weights matrix.
      recurrent_constraint: Constraint function applied to
        the `recurrent_kernel` weights matrix.
      bias_constraint: Constraint function applied to the bias vector.
      dropout: Float between 0 and 1.
        Fraction of the units to drop for
        the linear transformation of the inputs.
      recurrent_dropout: Float between 0 and 1.
        Fraction of the units to drop for
        the linear transformation of the recurrent state.
      implementation: Implementation mode, either 1 or 2.
        Mode 1 will structure its operations as a larger number of
        smaller dot products and additions, whereas mode 2 will
        batch them into fewer, larger operations. These modes will
        have different performance profiles on different hardware and
        for different applications.
    Call arguments:
      inputs: A 2D tensor.
      states: List of state tensors corresponding to the previous timestep.
      training: Python boolean indicating whether the layer should behave in
        training mode or in inference mode. Only relevant when `dropout` or
        `recurrent_dropout` is used.
    """

    def __init__(self, units, activation=tf.nn.tanh, recurrent_activation=
        tf.nn.sigmoid, use_bias=True, kernel_initializer='glorot_uniform',
        recurrent_initializer='orthogonal', bias_initializer='zeros',
        unit_forget_bias=True, kernel_regularizer=None,
        recurrent_regularizer=None, bias_regularizer=None,
        kernel_constraint=None, recurrent_constraint=None, bias_constraint=
        None, dropout=0.0, recurrent_dropout=0.0, implementation=2, **kwargs):
        super(LSTMCell, self).__init__(**kwargs)
        self.units = units
        self.activation = activation
        self.recurrent_activation = recurrent_activation
        self.use_bias = use_bias
        self.kernel_initializer = tf.keras.initializers.get(kernel_initializer)
        self.recurrent_initializer = tf.keras.initializers.get(
            recurrent_initializer)
        self.bias_initializer = tf.keras.initializers.get(bias_initializer)
        self.unit_forget_bias = unit_forget_bias
        self.kernel_regularizer = tf.keras.regularizers.get(kernel_regularizer)
        self.recurrent_regularizer = tf.keras.regularizers.get(
            recurrent_regularizer)
        self.bias_regularizer = tf.keras.regularizers.get(bias_regularizer)
        self.kernel_constraint = tf.keras.constraints.get(kernel_constraint)
        self.recurrent_constraint = tf.keras.constraints.get(
            recurrent_constraint)
        self.bias_constraint = tf.keras.constraints.get(bias_constraint)
        self.dropout = min(1.0, max(0.0, dropout))
        self.recurrent_dropout = min(1.0, max(0.0, recurrent_dropout))
        self.implementation = implementation
        self.state_size = data_structures.NoDependency([self.units, self.units]
            )
        self.output_size = self.units

    @tf_utils.shape_type_conversion
    def build(self, input_shape):
        input_dim = input_shape[-1]
        self.kernel = self.add_weight(shape=(input_dim, self.units * 4),
            name='kernel', initializer=self.kernel_initializer, regularizer
            =self.kernel_regularizer, constraint=self.kernel_constraint)
        self.recurrent_kernel = self.add_weight(shape=(self.units, self.
            units * 4), name='recurrent_kernel', initializer=self.
            recurrent_initializer, regularizer=self.recurrent_regularizer,
            constraint=self.recurrent_constraint)

        def exp_random_init(shape, dtype):
            rand = tf.random.uniform(shape, minval=1.0, maxval=6.0, dtype=dtype
                )
            return tf.exp(rand)

        def another_init(shape, dtype):
            return tf.random.uniform(shape, minval=0.0, maxval=1.0, dtype=dtype
                )
        self.periods = self.add_weight(shape=[self.units], name=
            'timegate_period', initializer=exp_random_init, constraint=tf.
            nn.relu)
        self.phases = self.add_weight(shape=[self.units], name=
            'timegate_phase', initializer=another_init, constraint=None)
        self.r_open = tf.constant(0.05 * np.ones([self.units], dtype=np.
            float32))
        self.leak = 0.001
        if self.use_bias:
            if self.unit_forget_bias:

                def bias_initializer(_, *args, **kwargs):
                    return K.concatenate([self.bias_initializer((self.units
                        ,), *args, **kwargs), tf.keras.initializers.Ones()(
                        (self.units,), *args, **kwargs), self.
                        bias_initializer((self.units * 2,), *args, **kwargs)])
            else:
                bias_initializer = self.bias_initializer
            self.bias = self.add_weight(shape=(self.units * 4,), name=
                'bias', initializer=bias_initializer, regularizer=self.
                bias_regularizer, constraint=self.bias_constraint)
        else:
            self.bias = None
        self.built = True

    def _compute_carry_and_output(self, x, h_tm1, c_tm1):
        """Computes carry and output using split kernels."""
        x_i, x_f, x_c, x_o = x
        h_tm1_i, h_tm1_f, h_tm1_c, h_tm1_o = h_tm1
        i = self.recurrent_activation(x_i + K.dot(h_tm1_i, self.
            recurrent_kernel[:, :self.units]))
        f = self.recurrent_activation(x_f + K.dot(h_tm1_f, self.
            recurrent_kernel[:, self.units:self.units * 2]))
        c = f * c_tm1 + i * self.activation(x_c + K.dot(h_tm1_c, self.
            recurrent_kernel[:, self.units * 2:self.units * 3]))
        o = self.recurrent_activation(x_o + K.dot(h_tm1_o, self.
            recurrent_kernel[:, self.units * 3:]))
        return c, o

    def _compute_carry_and_output_fused(self, z, c_tm1):
        """Computes carry and output using fused kernels."""
        z0, z1, z2, z3 = z
        i = self.recurrent_activation(z0)
        f = self.recurrent_activation(z1)
        c = f * c_tm1 + i * self.activation(z2)
        o = self.recurrent_activation(z3)
        return c, o

    def call(self, inputs, states, training=None):
        inputs, inp_time = inputs
        h_tm1 = states[0]
        c_tm1 = states[1]
        dp_mask = self.get_dropout_mask_for_cell(inputs, training, count=4)
        rec_dp_mask = self.get_recurrent_dropout_mask_for_cell(h_tm1,
            training, count=4)
        if self.implementation == 1:
            if 0 < self.dropout < 1.0:
                inputs_i = inputs * dp_mask[0]
                inputs_f = inputs * dp_mask[1]
                inputs_c = inputs * dp_mask[2]
                inputs_o = inputs * dp_mask[3]
            else:
                inputs_i = inputs
                inputs_f = inputs
                inputs_c = inputs
                inputs_o = inputs
            k_i, k_f, k_c, k_o = array_ops.split(self.kernel,
                num_or_size_splits=4, axis=1)
            x_i = K.dot(inputs_i, k_i)
            x_f = K.dot(inputs_f, k_f)
            x_c = K.dot(inputs_c, k_c)
            x_o = K.dot(inputs_o, k_o)
            if self.use_bias:
                b_i, b_f, b_c, b_o = array_ops.split(self.bias,
                    num_or_size_splits=4, axis=0)
                x_i = K.bias_add(x_i, b_i)
                x_f = K.bias_add(x_f, b_f)
                x_c = K.bias_add(x_c, b_c)
                x_o = K.bias_add(x_o, b_o)
            if 0 < self.recurrent_dropout < 1.0:
                h_tm1_i = h_tm1 * rec_dp_mask[0]
                h_tm1_f = h_tm1 * rec_dp_mask[1]
                h_tm1_c = h_tm1 * rec_dp_mask[2]
                h_tm1_o = h_tm1 * rec_dp_mask[3]
            else:
                h_tm1_i = h_tm1
                h_tm1_f = h_tm1
                h_tm1_c = h_tm1
                h_tm1_o = h_tm1
            x = x_i, x_f, x_c, x_o
            h_tm1 = h_tm1_i, h_tm1_f, h_tm1_c, h_tm1_o
            c, o = self._compute_carry_and_output(x, h_tm1, c_tm1)
        else:
            if 0.0 < self.dropout < 1.0:
                inputs *= dp_mask[0]
            z = K.dot(inputs, self.kernel)
            if 0.0 < self.recurrent_dropout < 1.0:
                h_tm1 *= rec_dp_mask[0]
            z += K.dot(h_tm1, self.recurrent_kernel)
            if self.use_bias:
                z = K.bias_add(z, self.bias)
            z = array_ops.split(z, num_or_size_splits=4, axis=1)
            c, o = self._compute_carry_and_output_fused(z, c_tm1)
        h = o * self.activation(c)
        aux = tf.math.mod(inp_time[:, tf.newaxis] - self.phases, self.periods
            ) / self.periods
        p1 = tf.cast(tf.less(aux, 0.5 * self.r_open), tf.float32
            ) * 2 * aux / self.r_open
        p2 = tf.cast(tf.logical_and(tf.greater_equal(aux, 0.5 * self.r_open
            ), tf.less(aux, self.r_open)), tf.float32) * (2 - 2 * aux /
            self.r_open)
        p3 = tf.cast(tf.greater_equal(aux, self.r_open), tf.float32
            ) * self.leak * aux
        time_gate = p1 + p2 + p3
        h = time_gate * h + (1 - time_gate) * h_tm1
        c = time_gate * c + (1 - time_gate) * c_tm1
        return h, [h, c]

    def get_config(self):
        config = {'units': self.units, 'activation': tf.keras.activations.
            serialize(self.activation), 'recurrent_activation': tf.keras.
            activations.serialize(self.recurrent_activation), 'use_bias':
            self.use_bias, 'kernel_initializer': tf.keras.initializers.
            serialize(self.kernel_initializer), 'recurrent_initializer': tf
            .keras.initializers.serialize(self.recurrent_initializer),
            'bias_initializer': tf.keras.initializers.serialize(self.
            bias_initializer), 'unit_forget_bias': self.unit_forget_bias,
            'kernel_regularizer': tf.keras.regularizers.serialize(self.
            kernel_regularizer), 'recurrent_regularizer': tf.keras.
            regularizers.serialize(self.recurrent_regularizer),
            'bias_regularizer': tf.keras.regularizers.serialize(self.
            bias_regularizer), 'kernel_constraint': tf.keras.constraints.
            serialize(self.kernel_constraint), 'recurrent_constraint': tf.
            keras.constraints.serialize(self.recurrent_constraint),
            'bias_constraint': tf.keras.constraints.serialize(self.
            bias_constraint), 'dropout': self.dropout, 'recurrent_dropout':
            self.recurrent_dropout, 'implementation': self.implementation}
        base_config = super(PhasedLSTMCell, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

    def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
        return list(_generate_zero_filled_state_for_cell(self, inputs,
            batch_size, dtype))


<mask token>

<mask token>


@register_keras_custom_object
class GroupNormalization(tf.keras.layers.Layer):
    <mask token>
    <mask token>

    def build(self, input_shape):
        self._check_if_input_shape_is_none(input_shape)
        self._set_number_of_groups_for_instance_norm(input_shape)
        self._check_size_of_dimensions(input_shape)
        self._create_input_spec(input_shape)
        self._add_gamma_weight(input_shape)
        self._add_beta_weight(input_shape)
        self.built = True
        super(GroupNormalization, self).build(input_shape)

    def call(self, inputs):
        input_shape = tf.keras.backend.int_shape(inputs)
        tensor_input_shape = tf.shape(inputs)
        reshaped_inputs, group_shape = self._reshape_into_groups(inputs,
            input_shape, tensor_input_shape)
        normalized_inputs = self._apply_normalization(reshaped_inputs,
            input_shape)
        outputs = tf.reshape(normalized_inputs, tensor_input_shape)
        return outputs

    def get_config(self):
        config = {'groups': self.groups, 'axis': self.axis, 'epsilon': self
            .epsilon, 'center': self.center, 'scale': self.scale,
            'beta_initializer': tf.keras.initializers.serialize(self.
            beta_initializer), 'gamma_initializer': tf.keras.initializers.
            serialize(self.gamma_initializer), 'beta_regularizer': tf.keras
            .regularizers.serialize(self.beta_regularizer),
            'gamma_regularizer': tf.keras.regularizers.serialize(self.
            gamma_regularizer), 'beta_constraint': tf.keras.constraints.
            serialize(self.beta_constraint), 'gamma_constraint': tf.keras.
            constraints.serialize(self.gamma_constraint)}
        base_config = super(GroupNormalization, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

    def compute_output_shape(self, input_shape):
        return input_shape

    def _reshape_into_groups(self, inputs, input_shape, tensor_input_shape):
        group_shape = [tensor_input_shape[i] for i in range(len(input_shape))]
        group_shape[self.axis] = input_shape[self.axis] // self.groups
        group_shape.insert(1, self.groups)
        group_shape = tf.stack(group_shape)
        reshaped_inputs = tf.reshape(inputs, group_shape)
        return reshaped_inputs, group_shape

    def _apply_normalization(self, reshaped_inputs, input_shape):
        group_shape = tf.keras.backend.int_shape(reshaped_inputs)
        group_reduction_axes = list(range(len(group_shape)))
        mean, variance = tf.nn.moments(reshaped_inputs,
            group_reduction_axes[2:], keepdims=True)
        gamma, beta = self._get_reshaped_weights(input_shape)
        normalized_inputs = tf.nn.batch_normalization(reshaped_inputs, mean
            =mean, variance=variance, scale=gamma, offset=beta,
            variance_epsilon=self.epsilon)
        return normalized_inputs

    def _get_reshaped_weights(self, input_shape):
        broadcast_shape = self._create_broadcast_shape(input_shape)
        gamma = None
        beta = None
        if self.scale:
            gamma = tf.reshape(self.gamma, broadcast_shape)
        if self.center:
            beta = tf.reshape(self.beta, broadcast_shape)
        return gamma, beta

    def _check_if_input_shape_is_none(self, input_shape):
        dim = input_shape[self.axis]
        if dim is None:
            raise ValueError('Axis ' + str(self.axis) +
                ' of input tensor should have a defined dimension but the layer received an input with shape '
                 + str(input_shape) + '.')

    def _set_number_of_groups_for_instance_norm(self, input_shape):
        dim = input_shape[self.axis]
        if self.groups == -1:
            self.groups = dim

    def _check_size_of_dimensions(self, input_shape):
        dim = input_shape[self.axis]
        if dim < self.groups:
            raise ValueError('Number of groups (' + str(self.groups) +
                ') cannot be more than the number of channels (' + str(dim) +
                ').')
        if dim % self.groups != 0:
            raise ValueError('Number of groups (' + str(self.groups) +
                ') must be a multiple of the number of channels (' + str(
                dim) + ').')

    def _check_axis(self):
        if self.axis == 0:
            raise ValueError(
                'You are trying to normalize your batch axis. Do you want to use tf.layer.batch_normalization instead'
                )
    <mask token>
    <mask token>

    def _add_beta_weight(self, input_shape):
        dim = input_shape[self.axis]
        shape = dim,
        if self.center:
            self.beta = self.add_weight(shape=shape, name='beta',
                initializer=self.beta_initializer, regularizer=self.
                beta_regularizer, constraint=self.beta_constraint)
        else:
            self.beta = None

    def _create_broadcast_shape(self, input_shape):
        broadcast_shape = [1] * len(input_shape)
        broadcast_shape[self.axis] = input_shape[self.axis] // self.groups
        broadcast_shape.insert(1, self.groups)
        return broadcast_shape


@register_keras_custom_object
class LayerNormalization(GroupNormalization):
    """Layer normalization layer.

    Layer Normalization is an specific case of ```GroupNormalization```since it
    normalizes all features of a layer. The Groupsize is 1.
    Empirically, its accuracy is more stable than batch norm in a wide
    range of small batch sizes, if learning rate is adjusted linearly
    with batch sizes.

    Arguments
        axis: Integer, the axis that should be normalized.
        epsilon: Small float added to variance to avoid dividing by zero.
        center: If True, add offset of `beta` to normalized tensor.
            If False, `beta` is ignored.
        scale: If True, multiply by `gamma`.
            If False, `gamma` is not used.
        beta_initializer: Initializer for the beta weight.
        gamma_initializer: Initializer for the gamma weight.
        beta_regularizer: Optional regularizer for the beta weight.
        gamma_regularizer: Optional regularizer for the gamma weight.
        beta_constraint: Optional constraint for the beta weight.
        gamma_constraint: Optional constraint for the gamma weight.

    Input shape
        Arbitrary. Use the keyword argument `input_shape`
        (tuple of integers, does not include the samples axis)
        when using this layer as the first layer in a model.

    Output shape
        Same shape as input.

    References
        - [Layer Normalization](https://arxiv.org/abs/1607.06450)
    """

    def __init__(self, **kwargs):
        if 'groups' in kwargs:
            logging.warning('The given value for groups will be overwritten.')
        kwargs['groups'] = 1
        super(LayerNormalization, self).__init__(**kwargs)


@register_keras_custom_object
class InstanceNormalization(GroupNormalization):
    """Instance normalization layer.

    Instance Normalization is an specific case of ```GroupNormalization```since
    it normalizes all features of one channel. The Groupsize is equal to the
    channel size. Empirically, its accuracy is more stable than batch norm in a
    wide range of small batch sizes, if learning rate is adjusted linearly
    with batch sizes.

    Arguments
        axis: Integer, the axis that should be normalized.
        epsilon: Small float added to variance to avoid dividing by zero.
        center: If True, add offset of `beta` to normalized tensor.
            If False, `beta` is ignored.
        scale: If True, multiply by `gamma`.
            If False, `gamma` is not used.
        beta_initializer: Initializer for the beta weight.
        gamma_initializer: Initializer for the gamma weight.
        beta_regularizer: Optional regularizer for the beta weight.
        gamma_regularizer: Optional regularizer for the gamma weight.
        beta_constraint: Optional constraint for the beta weight.
        gamma_constraint: Optional constraint for the gamma weight.

    Input shape
        Arbitrary. Use the keyword argument `input_shape`
        (tuple of integers, does not include the samples axis)
        when using this layer as the first layer in a model.

    Output shape
        Same shape as input.

    References
        - [Instance Normalization: The Missing Ingredient for Fast Stylization]
        (https://arxiv.org/abs/1607.08022)
    """

    def __init__(self, **kwargs):
        if 'groups' in kwargs:
            logging.warning('The given value for groups will be overwritten.')
        kwargs['groups'] = -1
        super(InstanceNormalization, self).__init__(**kwargs)


@register_keras_custom_object
class ResidualBlock(layers.Layer):

    def __init__(self, units, filter_size, activation):
        super(ResidualBlock, self).__init__()
        self.l1 = layers.Conv2D(units, filter_size, padding='same')
        self.b1 = layers.BatchNormalization()
        self.l2 = layers.Conv2D(units, filter_size, padding='same')
        self.b2 = layers.BatchNormalization()
        self.units = units
        self.activation = activation
        self.shortcut = None

    def build(self, input_shape):
        if input_shape[-1] != self.units:
            self.shortcut = layers.Conv2D(self.units, 1, padding='same')
            self.shortcut.build(input_shape)

    def call(self, inputs):
        interm = self.activation(self.b1(self.l1(inputs)))
        interm = self.b2(self.l2(interm))
        if self.shortcut:
            inputs = self.shortcut(inputs)
        return self.activation(inputs + interm)


@register_keras_custom_object
class PhasedLSTMCell(DropoutRNNCellMixin, layers.Layer):
    """Cell class for the LSTM layer.
    Arguments:
      units: Positive integer, dimensionality of the output space.
      activation: Activation function to use.
        Default: hyperbolic tangent (`tanh`).
        If you pass `None`, no activation is applied
        (ie. "linear" activation: `a(x) = x`).
      recurrent_activation: Activation function to use
        for the recurrent step.
        Default: hard sigmoid (`hard_sigmoid`).
        If you pass `None`, no activation is applied
        (ie. "linear" activation: `a(x) = x`).
      use_bias: Boolean, whether the layer uses a bias vector.
      kernel_initializer: Initializer for the `kernel` weights matrix,
        used for the linear transformation of the inputs.
      recurrent_initializer: Initializer for the `recurrent_kernel`
        weights matrix,
        used for the linear transformation of the recurrent state.
      bias_initializer: Initializer for the bias vector.
      unit_forget_bias: Boolean.
        If True, add 1 to the bias of the forget gate at initialization.
        Setting it to true will also force `bias_initializer="zeros"`.
        This is recommended in [Jozefowicz et
          al.](http://www.jmlr.org/proceedings/papers/v37/jozefowicz15.pdf)
      kernel_regularizer: Regularizer function applied to
        the `kernel` weights matrix.
      recurrent_regularizer: Regularizer function applied to
        the `recurrent_kernel` weights matrix.
      bias_regularizer: Regularizer function applied to the bias vector.
      kernel_constraint: Constraint function applied to
        the `kernel` weights matrix.
      recurrent_constraint: Constraint function applied to
        the `recurrent_kernel` weights matrix.
      bias_constraint: Constraint function applied to the bias vector.
      dropout: Float between 0 and 1.
        Fraction of the units to drop for
        the linear transformation of the inputs.
      recurrent_dropout: Float between 0 and 1.
        Fraction of the units to drop for
        the linear transformation of the recurrent state.
      implementation: Implementation mode, either 1 or 2.
        Mode 1 will structure its operations as a larger number of
        smaller dot products and additions, whereas mode 2 will
        batch them into fewer, larger operations. These modes will
        have different performance profiles on different hardware and
        for different applications.
    Call arguments:
      inputs: A 2D tensor.
      states: List of state tensors corresponding to the previous timestep.
      training: Python boolean indicating whether the layer should behave in
        training mode or in inference mode. Only relevant when `dropout` or
        `recurrent_dropout` is used.
    """

    def __init__(self, units, activation=tf.nn.tanh, recurrent_activation=
        tf.nn.sigmoid, use_bias=True, kernel_initializer='glorot_uniform',
        recurrent_initializer='orthogonal', bias_initializer='zeros',
        unit_forget_bias=True, kernel_regularizer=None,
        recurrent_regularizer=None, bias_regularizer=None,
        kernel_constraint=None, recurrent_constraint=None, bias_constraint=
        None, dropout=0.0, recurrent_dropout=0.0, implementation=2, **kwargs):
        super(LSTMCell, self).__init__(**kwargs)
        self.units = units
        self.activation = activation
        self.recurrent_activation = recurrent_activation
        self.use_bias = use_bias
        self.kernel_initializer = tf.keras.initializers.get(kernel_initializer)
        self.recurrent_initializer = tf.keras.initializers.get(
            recurrent_initializer)
        self.bias_initializer = tf.keras.initializers.get(bias_initializer)
        self.unit_forget_bias = unit_forget_bias
        self.kernel_regularizer = tf.keras.regularizers.get(kernel_regularizer)
        self.recurrent_regularizer = tf.keras.regularizers.get(
            recurrent_regularizer)
        self.bias_regularizer = tf.keras.regularizers.get(bias_regularizer)
        self.kernel_constraint = tf.keras.constraints.get(kernel_constraint)
        self.recurrent_constraint = tf.keras.constraints.get(
            recurrent_constraint)
        self.bias_constraint = tf.keras.constraints.get(bias_constraint)
        self.dropout = min(1.0, max(0.0, dropout))
        self.recurrent_dropout = min(1.0, max(0.0, recurrent_dropout))
        self.implementation = implementation
        self.state_size = data_structures.NoDependency([self.units, self.units]
            )
        self.output_size = self.units

    @tf_utils.shape_type_conversion
    def build(self, input_shape):
        input_dim = input_shape[-1]
        self.kernel = self.add_weight(shape=(input_dim, self.units * 4),
            name='kernel', initializer=self.kernel_initializer, regularizer
            =self.kernel_regularizer, constraint=self.kernel_constraint)
        self.recurrent_kernel = self.add_weight(shape=(self.units, self.
            units * 4), name='recurrent_kernel', initializer=self.
            recurrent_initializer, regularizer=self.recurrent_regularizer,
            constraint=self.recurrent_constraint)

        def exp_random_init(shape, dtype):
            rand = tf.random.uniform(shape, minval=1.0, maxval=6.0, dtype=dtype
                )
            return tf.exp(rand)

        def another_init(shape, dtype):
            return tf.random.uniform(shape, minval=0.0, maxval=1.0, dtype=dtype
                )
        self.periods = self.add_weight(shape=[self.units], name=
            'timegate_period', initializer=exp_random_init, constraint=tf.
            nn.relu)
        self.phases = self.add_weight(shape=[self.units], name=
            'timegate_phase', initializer=another_init, constraint=None)
        self.r_open = tf.constant(0.05 * np.ones([self.units], dtype=np.
            float32))
        self.leak = 0.001
        if self.use_bias:
            if self.unit_forget_bias:

                def bias_initializer(_, *args, **kwargs):
                    return K.concatenate([self.bias_initializer((self.units
                        ,), *args, **kwargs), tf.keras.initializers.Ones()(
                        (self.units,), *args, **kwargs), self.
                        bias_initializer((self.units * 2,), *args, **kwargs)])
            else:
                bias_initializer = self.bias_initializer
            self.bias = self.add_weight(shape=(self.units * 4,), name=
                'bias', initializer=bias_initializer, regularizer=self.
                bias_regularizer, constraint=self.bias_constraint)
        else:
            self.bias = None
        self.built = True

    def _compute_carry_and_output(self, x, h_tm1, c_tm1):
        """Computes carry and output using split kernels."""
        x_i, x_f, x_c, x_o = x
        h_tm1_i, h_tm1_f, h_tm1_c, h_tm1_o = h_tm1
        i = self.recurrent_activation(x_i + K.dot(h_tm1_i, self.
            recurrent_kernel[:, :self.units]))
        f = self.recurrent_activation(x_f + K.dot(h_tm1_f, self.
            recurrent_kernel[:, self.units:self.units * 2]))
        c = f * c_tm1 + i * self.activation(x_c + K.dot(h_tm1_c, self.
            recurrent_kernel[:, self.units * 2:self.units * 3]))
        o = self.recurrent_activation(x_o + K.dot(h_tm1_o, self.
            recurrent_kernel[:, self.units * 3:]))
        return c, o

    def _compute_carry_and_output_fused(self, z, c_tm1):
        """Computes carry and output using fused kernels."""
        z0, z1, z2, z3 = z
        i = self.recurrent_activation(z0)
        f = self.recurrent_activation(z1)
        c = f * c_tm1 + i * self.activation(z2)
        o = self.recurrent_activation(z3)
        return c, o

    def call(self, inputs, states, training=None):
        inputs, inp_time = inputs
        h_tm1 = states[0]
        c_tm1 = states[1]
        dp_mask = self.get_dropout_mask_for_cell(inputs, training, count=4)
        rec_dp_mask = self.get_recurrent_dropout_mask_for_cell(h_tm1,
            training, count=4)
        if self.implementation == 1:
            if 0 < self.dropout < 1.0:
                inputs_i = inputs * dp_mask[0]
                inputs_f = inputs * dp_mask[1]
                inputs_c = inputs * dp_mask[2]
                inputs_o = inputs * dp_mask[3]
            else:
                inputs_i = inputs
                inputs_f = inputs
                inputs_c = inputs
                inputs_o = inputs
            k_i, k_f, k_c, k_o = array_ops.split(self.kernel,
                num_or_size_splits=4, axis=1)
            x_i = K.dot(inputs_i, k_i)
            x_f = K.dot(inputs_f, k_f)
            x_c = K.dot(inputs_c, k_c)
            x_o = K.dot(inputs_o, k_o)
            if self.use_bias:
                b_i, b_f, b_c, b_o = array_ops.split(self.bias,
                    num_or_size_splits=4, axis=0)
                x_i = K.bias_add(x_i, b_i)
                x_f = K.bias_add(x_f, b_f)
                x_c = K.bias_add(x_c, b_c)
                x_o = K.bias_add(x_o, b_o)
            if 0 < self.recurrent_dropout < 1.0:
                h_tm1_i = h_tm1 * rec_dp_mask[0]
                h_tm1_f = h_tm1 * rec_dp_mask[1]
                h_tm1_c = h_tm1 * rec_dp_mask[2]
                h_tm1_o = h_tm1 * rec_dp_mask[3]
            else:
                h_tm1_i = h_tm1
                h_tm1_f = h_tm1
                h_tm1_c = h_tm1
                h_tm1_o = h_tm1
            x = x_i, x_f, x_c, x_o
            h_tm1 = h_tm1_i, h_tm1_f, h_tm1_c, h_tm1_o
            c, o = self._compute_carry_and_output(x, h_tm1, c_tm1)
        else:
            if 0.0 < self.dropout < 1.0:
                inputs *= dp_mask[0]
            z = K.dot(inputs, self.kernel)
            if 0.0 < self.recurrent_dropout < 1.0:
                h_tm1 *= rec_dp_mask[0]
            z += K.dot(h_tm1, self.recurrent_kernel)
            if self.use_bias:
                z = K.bias_add(z, self.bias)
            z = array_ops.split(z, num_or_size_splits=4, axis=1)
            c, o = self._compute_carry_and_output_fused(z, c_tm1)
        h = o * self.activation(c)
        aux = tf.math.mod(inp_time[:, tf.newaxis] - self.phases, self.periods
            ) / self.periods
        p1 = tf.cast(tf.less(aux, 0.5 * self.r_open), tf.float32
            ) * 2 * aux / self.r_open
        p2 = tf.cast(tf.logical_and(tf.greater_equal(aux, 0.5 * self.r_open
            ), tf.less(aux, self.r_open)), tf.float32) * (2 - 2 * aux /
            self.r_open)
        p3 = tf.cast(tf.greater_equal(aux, self.r_open), tf.float32
            ) * self.leak * aux
        time_gate = p1 + p2 + p3
        h = time_gate * h + (1 - time_gate) * h_tm1
        c = time_gate * c + (1 - time_gate) * c_tm1
        return h, [h, c]

    def get_config(self):
        config = {'units': self.units, 'activation': tf.keras.activations.
            serialize(self.activation), 'recurrent_activation': tf.keras.
            activations.serialize(self.recurrent_activation), 'use_bias':
            self.use_bias, 'kernel_initializer': tf.keras.initializers.
            serialize(self.kernel_initializer), 'recurrent_initializer': tf
            .keras.initializers.serialize(self.recurrent_initializer),
            'bias_initializer': tf.keras.initializers.serialize(self.
            bias_initializer), 'unit_forget_bias': self.unit_forget_bias,
            'kernel_regularizer': tf.keras.regularizers.serialize(self.
            kernel_regularizer), 'recurrent_regularizer': tf.keras.
            regularizers.serialize(self.recurrent_regularizer),
            'bias_regularizer': tf.keras.regularizers.serialize(self.
            bias_regularizer), 'kernel_constraint': tf.keras.constraints.
            serialize(self.kernel_constraint), 'recurrent_constraint': tf.
            keras.constraints.serialize(self.recurrent_constraint),
            'bias_constraint': tf.keras.constraints.serialize(self.
            bias_constraint), 'dropout': self.dropout, 'recurrent_dropout':
            self.recurrent_dropout, 'implementation': self.implementation}
        base_config = super(PhasedLSTMCell, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

    def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
        return list(_generate_zero_filled_state_for_cell(self, inputs,
            batch_size, dtype))


<mask token>

import logging

import tensorflow as tf
import tensorflow.keras.layers as layers
from tensorflow.keras.utils import get_custom_objects
from tensorflow.python.training.tracking import data_structures
from tensorflow.python.keras.layers.recurrent import DropoutRNNCellMixin
from tensorflow.python.keras.utils import tf_utils
from tensorflow.keras import backend as K
from tensorflow.python.util import nest
from tensorflow.python.ops import array_ops
from tensorflow.python.framework import tensor_shape
import numpy as np


################################################################################
# I STOLE THIS FROM THE TF-ADDONS REPO BECAUSE THEY DON'T PROVIDE A MAC VERSION
################################################################################
def register_keras_custom_object(cls):
    get_custom_objects()[cls.__name__] = cls
    return cls


@register_keras_custom_object
class GroupNormalization(tf.keras.layers.Layer):
    """Group normalization layer.

    Group Normalization divides the channels into groups and computes
    within each group the mean and variance for normalization.
    Empirically, its accuracy is more stable than batch norm in a wide
    range of small batch sizes, if learning rate is adjusted linearly
    with batch sizes.

    Relation to Layer Normalization:
    If the number of groups is set to 1, then this operation becomes identical
    to Layer Normalization.

    Relation to Instance Normalization:
    If the number of groups is set to the
    input dimension (number of groups is equal
    to number of channels), then this operation becomes
    identical to Instance Normalization.

    Arguments
        groups: Integer, the number of groups for Group Normalization.
            Can be in the range [1, N] where N is the input dimension.
            The input dimension must be divisible by the number of groups.
        axis: Integer, the axis that should be normalized.
        epsilon: Small float added to variance to avoid dividing by zero.
        center: If True, add offset of `beta` to normalized tensor.
            If False, `beta` is ignored.
        scale: If True, multiply by `gamma`.
            If False, `gamma` is not used.
        beta_initializer: Initializer for the beta weight.
        gamma_initializer: Initializer for the gamma weight.
        beta_regularizer: Optional regularizer for the beta weight.
        gamma_regularizer: Optional regularizer for the gamma weight.
        beta_constraint: Optional constraint for the beta weight.
        gamma_constraint: Optional constraint for the gamma weight.

    Input shape
        Arbitrary. Use the keyword argument `input_shape`
        (tuple of integers, does not include the samples axis)
        when using this layer as the first layer in a model.

    Output shape
        Same shape as input.
    References
        - [Group Normalization](https://arxiv.org/abs/1803.08494)
    """

    def __init__(self,
                 groups=2,
                 axis=-1,
                 epsilon=1e-5,
                 center=True,
                 scale=True,
                 beta_initializer='zeros',
                 gamma_initializer='ones',
                 beta_regularizer=None,
                 gamma_regularizer=None,
                 beta_constraint=None,
                 gamma_constraint=None,
                 **kwargs):
        super(GroupNormalization, self).__init__(**kwargs)
        self.supports_masking = True
        self.groups = groups
        self.axis = axis
        self.epsilon = epsilon
        self.center = center
        self.scale = scale
        self.beta_initializer = tf.keras.initializers.get(beta_initializer)
        self.gamma_initializer = tf.keras.initializers.get(gamma_initializer)
        self.beta_regularizer = tf.keras.regularizers.get(beta_regularizer)
        self.gamma_regularizer = tf.keras.regularizers.get(gamma_regularizer)
        self.beta_constraint = tf.keras.constraints.get(beta_constraint)
        self.gamma_constraint = tf.keras.constraints.get(gamma_constraint)
        self._check_axis()

    def build(self, input_shape):

        self._check_if_input_shape_is_none(input_shape)
        self._set_number_of_groups_for_instance_norm(input_shape)
        self._check_size_of_dimensions(input_shape)
        self._create_input_spec(input_shape)

        self._add_gamma_weight(input_shape)
        self._add_beta_weight(input_shape)
        self.built = True
        super(GroupNormalization, self).build(input_shape)

    def call(self, inputs):

        input_shape = tf.keras.backend.int_shape(inputs)
        tensor_input_shape = tf.shape(inputs)

        reshaped_inputs, group_shape = self._reshape_into_groups(
            inputs, input_shape, tensor_input_shape)

        normalized_inputs = self._apply_normalization(reshaped_inputs,
                                                      input_shape)

        outputs = tf.reshape(normalized_inputs, tensor_input_shape)

        return outputs

    def get_config(self):
        config = {
            'groups':
                self.groups,
            'axis':
                self.axis,
            'epsilon':
                self.epsilon,
            'center':
                self.center,
            'scale':
                self.scale,
            'beta_initializer':
                tf.keras.initializers.serialize(self.beta_initializer),
            'gamma_initializer':
                tf.keras.initializers.serialize(self.gamma_initializer),
            'beta_regularizer':
                tf.keras.regularizers.serialize(self.beta_regularizer),
            'gamma_regularizer':
                tf.keras.regularizers.serialize(self.gamma_regularizer),
            'beta_constraint':
                tf.keras.constraints.serialize(self.beta_constraint),
            'gamma_constraint':
                tf.keras.constraints.serialize(self.gamma_constraint)
        }
        base_config = super(GroupNormalization, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

    def compute_output_shape(self, input_shape):
        return input_shape

    def _reshape_into_groups(self, inputs, input_shape, tensor_input_shape):

        group_shape = [tensor_input_shape[i] for i in range(len(input_shape))]
        group_shape[self.axis] = input_shape[self.axis] // self.groups
        group_shape.insert(1, self.groups)
        group_shape = tf.stack(group_shape)
        reshaped_inputs = tf.reshape(inputs, group_shape)
        return reshaped_inputs, group_shape

    def _apply_normalization(self, reshaped_inputs, input_shape):

        group_shape = tf.keras.backend.int_shape(reshaped_inputs)
        group_reduction_axes = list(range(len(group_shape)))
        # Remember the ordering of the tensor is [batch, group , steps]. Jump
        # the first 2 to calculate the variance and the mean
        mean, variance = tf.nn.moments(
            reshaped_inputs, group_reduction_axes[2:], keepdims=True)

        gamma, beta = self._get_reshaped_weights(input_shape)
        normalized_inputs = tf.nn.batch_normalization(
            reshaped_inputs,
            mean=mean,
            variance=variance,
            scale=gamma,
            offset=beta,
            variance_epsilon=self.epsilon)
        return normalized_inputs

    def _get_reshaped_weights(self, input_shape):
        broadcast_shape = self._create_broadcast_shape(input_shape)
        gamma = None
        beta = None
        if self.scale:
            gamma = tf.reshape(self.gamma, broadcast_shape)

        if self.center:
            beta = tf.reshape(self.beta, broadcast_shape)
        return gamma, beta

    def _check_if_input_shape_is_none(self, input_shape):
        dim = input_shape[self.axis]
        if dim is None:
            raise ValueError('Axis ' + str(self.axis) + ' of '
                                                        'input tensor should have a defined dimension '
                                                        'but the layer received an input with shape ' +
                             str(input_shape) + '.')

    def _set_number_of_groups_for_instance_norm(self, input_shape):
        dim = input_shape[self.axis]

        if self.groups == -1:
            self.groups = dim

    def _check_size_of_dimensions(self, input_shape):

        dim = input_shape[self.axis]
        if dim < self.groups:
            raise ValueError(
                'Number of groups (' + str(self.groups) + ') cannot be '
                                                          'more than the number of channels (' + str(
                    dim) + ').')

        if dim % self.groups != 0:
            raise ValueError(
                'Number of groups (' + str(self.groups) + ') must be a '
                                                          'multiple of the number of channels (' + str(
                    dim) + ').')

    def _check_axis(self):

        if self.axis == 0:
            raise ValueError(
                "You are trying to normalize your batch axis. Do you want to "
                "use tf.layer.batch_normalization instead")

    def _create_input_spec(self, input_shape):

        dim = input_shape[self.axis]
        self.input_spec = tf.keras.layers.InputSpec(
            ndim=len(input_shape), axes={self.axis: dim})

    def _add_gamma_weight(self, input_shape):

        dim = input_shape[self.axis]
        shape = (dim,)

        if self.scale:
            self.gamma = self.add_weight(
                shape=shape,
                name='gamma',
                initializer=self.gamma_initializer,
                regularizer=self.gamma_regularizer,
                constraint=self.gamma_constraint)
        else:
            self.gamma = None

    def _add_beta_weight(self, input_shape):

        dim = input_shape[self.axis]
        shape = (dim,)

        if self.center:
            self.beta = self.add_weight(
                shape=shape,
                name='beta',
                initializer=self.beta_initializer,
                regularizer=self.beta_regularizer,
                constraint=self.beta_constraint)
        else:
            self.beta = None

    def _create_broadcast_shape(self, input_shape):
        broadcast_shape = [1] * len(input_shape)
        broadcast_shape[self.axis] = input_shape[self.axis] // self.groups
        broadcast_shape.insert(1, self.groups)
        return broadcast_shape


@register_keras_custom_object
class LayerNormalization(GroupNormalization):
    """Layer normalization layer.

    Layer Normalization is an specific case of ```GroupNormalization```since it
    normalizes all features of a layer. The Groupsize is 1.
    Empirically, its accuracy is more stable than batch norm in a wide
    range of small batch sizes, if learning rate is adjusted linearly
    with batch sizes.

    Arguments
        axis: Integer, the axis that should be normalized.
        epsilon: Small float added to variance to avoid dividing by zero.
        center: If True, add offset of `beta` to normalized tensor.
            If False, `beta` is ignored.
        scale: If True, multiply by `gamma`.
            If False, `gamma` is not used.
        beta_initializer: Initializer for the beta weight.
        gamma_initializer: Initializer for the gamma weight.
        beta_regularizer: Optional regularizer for the beta weight.
        gamma_regularizer: Optional regularizer for the gamma weight.
        beta_constraint: Optional constraint for the beta weight.
        gamma_constraint: Optional constraint for the gamma weight.

    Input shape
        Arbitrary. Use the keyword argument `input_shape`
        (tuple of integers, does not include the samples axis)
        when using this layer as the first layer in a model.

    Output shape
        Same shape as input.

    References
        - [Layer Normalization](https://arxiv.org/abs/1607.06450)
    """

    def __init__(self, **kwargs):
        if "groups" in kwargs:
            logging.warning("The given value for groups will be overwritten.")
        kwargs["groups"] = 1
        super(LayerNormalization, self).__init__(**kwargs)


@register_keras_custom_object
class InstanceNormalization(GroupNormalization):
    """Instance normalization layer.

    Instance Normalization is an specific case of ```GroupNormalization```since
    it normalizes all features of one channel. The Groupsize is equal to the
    channel size. Empirically, its accuracy is more stable than batch norm in a
    wide range of small batch sizes, if learning rate is adjusted linearly
    with batch sizes.

    Arguments
        axis: Integer, the axis that should be normalized.
        epsilon: Small float added to variance to avoid dividing by zero.
        center: If True, add offset of `beta` to normalized tensor.
            If False, `beta` is ignored.
        scale: If True, multiply by `gamma`.
            If False, `gamma` is not used.
        beta_initializer: Initializer for the beta weight.
        gamma_initializer: Initializer for the gamma weight.
        beta_regularizer: Optional regularizer for the beta weight.
        gamma_regularizer: Optional regularizer for the gamma weight.
        beta_constraint: Optional constraint for the beta weight.
        gamma_constraint: Optional constraint for the gamma weight.

    Input shape
        Arbitrary. Use the keyword argument `input_shape`
        (tuple of integers, does not include the samples axis)
        when using this layer as the first layer in a model.

    Output shape
        Same shape as input.

    References
        - [Instance Normalization: The Missing Ingredient for Fast Stylization]
        (https://arxiv.org/abs/1607.08022)
    """

    def __init__(self, **kwargs):
        if "groups" in kwargs:
            logging.warning("The given value for groups will be overwritten.")

        kwargs["groups"] = -1
        super(InstanceNormalization, self).__init__(**kwargs)


################################################################################
# END OF THIEVERY
################################################################################


@register_keras_custom_object
class ResidualBlock(layers.Layer):
    def __init__(self, units, filter_size, activation):
        super(ResidualBlock, self).__init__()
        self.l1 = layers.Conv2D(units, filter_size, padding="same")
        self.b1 = layers.BatchNormalization()
        self.l2 = layers.Conv2D(units, filter_size, padding="same")
        self.b2 = layers.BatchNormalization()
        self.units = units
        self.activation = activation
        self.shortcut = None

    def build(self, input_shape):
        if input_shape[-1] != self.units:
            self.shortcut = layers.Conv2D(self.units, 1, padding="same")
            self.shortcut.build(input_shape)

    def call(self, inputs):
        interm = self.activation(self.b1(self.l1(inputs)))
        interm = self.b2(self.l2(interm))
        if self.shortcut:
            inputs = self.shortcut(inputs)
        return self.activation(inputs + interm)


@register_keras_custom_object
class PhasedLSTMCell(DropoutRNNCellMixin, layers.Layer):
    """Cell class for the LSTM layer.
    Arguments:
      units: Positive integer, dimensionality of the output space.
      activation: Activation function to use.
        Default: hyperbolic tangent (`tanh`).
        If you pass `None`, no activation is applied
        (ie. "linear" activation: `a(x) = x`).
      recurrent_activation: Activation function to use
        for the recurrent step.
        Default: hard sigmoid (`hard_sigmoid`).
        If you pass `None`, no activation is applied
        (ie. "linear" activation: `a(x) = x`).
      use_bias: Boolean, whether the layer uses a bias vector.
      kernel_initializer: Initializer for the `kernel` weights matrix,
        used for the linear transformation of the inputs.
      recurrent_initializer: Initializer for the `recurrent_kernel`
        weights matrix,
        used for the linear transformation of the recurrent state.
      bias_initializer: Initializer for the bias vector.
      unit_forget_bias: Boolean.
        If True, add 1 to the bias of the forget gate at initialization.
        Setting it to true will also force `bias_initializer="zeros"`.
        This is recommended in [Jozefowicz et
          al.](http://www.jmlr.org/proceedings/papers/v37/jozefowicz15.pdf)
      kernel_regularizer: Regularizer function applied to
        the `kernel` weights matrix.
      recurrent_regularizer: Regularizer function applied to
        the `recurrent_kernel` weights matrix.
      bias_regularizer: Regularizer function applied to the bias vector.
      kernel_constraint: Constraint function applied to
        the `kernel` weights matrix.
      recurrent_constraint: Constraint function applied to
        the `recurrent_kernel` weights matrix.
      bias_constraint: Constraint function applied to the bias vector.
      dropout: Float between 0 and 1.
        Fraction of the units to drop for
        the linear transformation of the inputs.
      recurrent_dropout: Float between 0 and 1.
        Fraction of the units to drop for
        the linear transformation of the recurrent state.
      implementation: Implementation mode, either 1 or 2.
        Mode 1 will structure its operations as a larger number of
        smaller dot products and additions, whereas mode 2 will
        batch them into fewer, larger operations. These modes will
        have different performance profiles on different hardware and
        for different applications.
    Call arguments:
      inputs: A 2D tensor.
      states: List of state tensors corresponding to the previous timestep.
      training: Python boolean indicating whether the layer should behave in
        training mode or in inference mode. Only relevant when `dropout` or
        `recurrent_dropout` is used.
    """

    def __init__(self,
                 units,
                 activation=tf.nn.tanh,
                 recurrent_activation=tf.nn.sigmoid,
                 use_bias=True,
                 kernel_initializer='glorot_uniform',
                 recurrent_initializer='orthogonal',
                 bias_initializer='zeros',
                 unit_forget_bias=True,
                 kernel_regularizer=None,
                 recurrent_regularizer=None,
                 bias_regularizer=None,
                 kernel_constraint=None,
                 recurrent_constraint=None,
                 bias_constraint=None,
                 dropout=0.,
                 recurrent_dropout=0.,
                 implementation=2,
                 **kwargs):
        super(LSTMCell, self).__init__(**kwargs)
        self.units = units
        self.activation = activation
        self.recurrent_activation = recurrent_activation
        self.use_bias = use_bias

        self.kernel_initializer = tf.keras.initializers.get(kernel_initializer)
        self.recurrent_initializer = tf.keras.initializers.get(
            recurrent_initializer)
        self.bias_initializer = tf.keras.initializers.get(bias_initializer)
        self.unit_forget_bias = unit_forget_bias

        self.kernel_regularizer = tf.keras.regularizers.get(kernel_regularizer)
        self.recurrent_regularizer = tf.keras.regularizers.get(
            recurrent_regularizer)
        self.bias_regularizer = tf.keras.regularizers.get(bias_regularizer)

        self.kernel_constraint = tf.keras.constraints.get(kernel_constraint)
        self.recurrent_constraint = tf.keras.constraints.get(
            recurrent_constraint)
        self.bias_constraint = tf.keras.constraints.get(bias_constraint)

        self.dropout = min(1., max(0., dropout))
        self.recurrent_dropout = min(1., max(0., recurrent_dropout))
        self.implementation = implementation
        # tuple(_ListWrapper) was silently dropping list content in at least 2.7.10,
        # and fixed after 2.7.16. Converting the state_size to wrapper around
        # NoDependency(), so that the base_layer.__setattr__ will not convert it to
        # ListWrapper. Down the stream, self.states will be a list since it is
        # generated from nest.map_structure with list, and tuple(list) will work
        # properly.
        self.state_size = data_structures.NoDependency([self.units, self.units])
        self.output_size = self.units

    @tf_utils.shape_type_conversion
    def build(self, input_shape):
        input_dim = input_shape[-1]
        self.kernel = self.add_weight(
            shape=(input_dim, self.units * 4),
            name='kernel',
            initializer=self.kernel_initializer,
            regularizer=self.kernel_regularizer,
            constraint=self.kernel_constraint)
        self.recurrent_kernel = self.add_weight(
            shape=(self.units, self.units * 4),
            name='recurrent_kernel',
            initializer=self.recurrent_initializer,
            regularizer=self.recurrent_regularizer,
            constraint=self.recurrent_constraint)

        def exp_random_init(shape, dtype):
            rand = tf.random.uniform(shape, minval=1., maxval=6., dtype=dtype)
            return tf.exp(rand)

        def another_init(shape, dtype):
            return tf.random.uniform(shape, minval=0., maxval=1., dtype=dtype)

        self.periods = self.add_weight(
            shape=[self.units],
            name="timegate_period",
            initializer=exp_random_init,
            constraint=tf.nn.relu)
        self.phases = self.add_weight(
            shape=[self.units],
            name="timegate_phase",
            initializer=another_init,
            constraint=None)
        self.r_open = tf.constant(
            0.05 * np.ones([self.units], dtype=np.float32))
        self.leak = 0.001

        if self.use_bias:
            if self.unit_forget_bias:

                def bias_initializer(_, *args, **kwargs):
                    return K.concatenate([
                        self.bias_initializer((self.units,), *args, **kwargs),
                        tf.keras.initializers.Ones()((self.units,), *args,
                                                     **kwargs),
                        self.bias_initializer((self.units * 2,), *args,
                                              **kwargs),
                    ])
            else:
                bias_initializer = self.bias_initializer
            self.bias = self.add_weight(
                shape=(self.units * 4,),
                name='bias',
                initializer=bias_initializer,
                regularizer=self.bias_regularizer,
                constraint=self.bias_constraint)
        else:
            self.bias = None
        self.built = True

    def _compute_carry_and_output(self, x, h_tm1, c_tm1):
        """Computes carry and output using split kernels."""
        x_i, x_f, x_c, x_o = x
        h_tm1_i, h_tm1_f, h_tm1_c, h_tm1_o = h_tm1
        i = self.recurrent_activation(
            x_i + K.dot(h_tm1_i, self.recurrent_kernel[:, :self.units]))
        f = self.recurrent_activation(x_f + K.dot(
            h_tm1_f, self.recurrent_kernel[:, self.units:self.units * 2]))
        c = f * c_tm1 + i * self.activation(x_c + K.dot(
            h_tm1_c, self.recurrent_kernel[:, self.units * 2:self.units * 3]))
        o = self.recurrent_activation(
            x_o + K.dot(h_tm1_o, self.recurrent_kernel[:, self.units * 3:]))
        return c, o

    def _compute_carry_and_output_fused(self, z, c_tm1):
        """Computes carry and output using fused kernels."""
        z0, z1, z2, z3 = z
        i = self.recurrent_activation(z0)
        f = self.recurrent_activation(z1)
        c = f * c_tm1 + i * self.activation(z2)
        o = self.recurrent_activation(z3)
        return c, o

    def call(self, inputs, states, training=None):
        inputs, inp_time = inputs

        h_tm1 = states[0]  # previous memory state
        c_tm1 = states[1]  # previous carry state

        dp_mask = self.get_dropout_mask_for_cell(inputs, training, count=4)
        rec_dp_mask = self.get_recurrent_dropout_mask_for_cell(
            h_tm1, training, count=4)

        if self.implementation == 1:
            if 0 < self.dropout < 1.:
                inputs_i = inputs * dp_mask[0]
                inputs_f = inputs * dp_mask[1]
                inputs_c = inputs * dp_mask[2]
                inputs_o = inputs * dp_mask[3]
            else:
                inputs_i = inputs
                inputs_f = inputs
                inputs_c = inputs
                inputs_o = inputs
            k_i, k_f, k_c, k_o = array_ops.split(
                self.kernel, num_or_size_splits=4, axis=1)
            x_i = K.dot(inputs_i, k_i)
            x_f = K.dot(inputs_f, k_f)
            x_c = K.dot(inputs_c, k_c)
            x_o = K.dot(inputs_o, k_o)
            if self.use_bias:
                b_i, b_f, b_c, b_o = array_ops.split(
                    self.bias, num_or_size_splits=4, axis=0)
                x_i = K.bias_add(x_i, b_i)
                x_f = K.bias_add(x_f, b_f)
                x_c = K.bias_add(x_c, b_c)
                x_o = K.bias_add(x_o, b_o)

            if 0 < self.recurrent_dropout < 1.:
                h_tm1_i = h_tm1 * rec_dp_mask[0]
                h_tm1_f = h_tm1 * rec_dp_mask[1]
                h_tm1_c = h_tm1 * rec_dp_mask[2]
                h_tm1_o = h_tm1 * rec_dp_mask[3]
            else:
                h_tm1_i = h_tm1
                h_tm1_f = h_tm1
                h_tm1_c = h_tm1
                h_tm1_o = h_tm1
            x = (x_i, x_f, x_c, x_o)
            h_tm1 = (h_tm1_i, h_tm1_f, h_tm1_c, h_tm1_o)
            c, o = self._compute_carry_and_output(x, h_tm1, c_tm1)
        else:
            if 0. < self.dropout < 1.:
                inputs *= dp_mask[0]
            z = K.dot(inputs, self.kernel)
            if 0. < self.recurrent_dropout < 1.:
                h_tm1 *= rec_dp_mask[0]
            z += K.dot(h_tm1, self.recurrent_kernel)
            if self.use_bias:
                z = K.bias_add(z, self.bias)

            z = array_ops.split(z, num_or_size_splits=4, axis=1)
            c, o = self._compute_carry_and_output_fused(z, c_tm1)

        h = o * self.activation(c)

        # time gate stufffffff
        # time_gate = 1.  # use this for debugging (always open)
        # inp_time = 0.
        # inp_time is just batch. so we reshape to batch x 1
        # so that then broadcasting can do the rest to make everything
        # batch x features
        aux = tf.math.mod(inp_time[:, tf.newaxis] - self.phases,
                          self.periods) / self.periods

        p1 = tf.cast(tf.less(aux, 0.5 * self.r_open),
                     tf.float32) * 2 * aux / self.r_open
        p2 = tf.cast(tf.logical_and(tf.greater_equal(aux, 0.5 * self.r_open),
                                    tf.less(aux, self.r_open)), tf.float32) * (
                         2 - 2 * aux / self.r_open)
        p3 = tf.cast(tf.greater_equal(aux, self.r_open),
                     tf.float32) * self.leak * aux
        time_gate = p1 + p2 + p3

        h = time_gate * h + (1 - time_gate) * h_tm1
        c = time_gate * c + (1 - time_gate) * c_tm1
        return h, [h, c]

    def get_config(self):
        config = {
            'units':
                self.units,
            'activation':
                tf.keras.activations.serialize(self.activation),
            'recurrent_activation':
                tf.keras.activations.serialize(self.recurrent_activation),
            'use_bias':
                self.use_bias,
            'kernel_initializer':
                tf.keras.initializers.serialize(self.kernel_initializer),
            'recurrent_initializer':
                tf.keras.initializers.serialize(self.recurrent_initializer),
            'bias_initializer':
                tf.keras.initializers.serialize(self.bias_initializer),
            'unit_forget_bias':
                self.unit_forget_bias,
            'kernel_regularizer':
                tf.keras.regularizers.serialize(self.kernel_regularizer),
            'recurrent_regularizer':
                tf.keras.regularizers.serialize(self.recurrent_regularizer),
            'bias_regularizer':
                tf.keras.regularizers.serialize(self.bias_regularizer),
            'kernel_constraint':
                tf.keras.constraints.serialize(self.kernel_constraint),
            'recurrent_constraint':
                tf.keras.constraints.serialize(self.recurrent_constraint),
            'bias_constraint':
                tf.keras.constraints.serialize(self.bias_constraint),
            'dropout':
                self.dropout,
            'recurrent_dropout':
                self.recurrent_dropout,
            'implementation':
                self.implementation
        }
        base_config = super(PhasedLSTMCell, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

    def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
        return list(_generate_zero_filled_state_for_cell(
            self, inputs, batch_size, dtype))


def _generate_zero_filled_state_for_cell(cell, inputs, batch_size, dtype):
    if inputs is not None:
        batch_size = array_ops.shape(inputs)[0]
        dtype = inputs.dtype
    return _generate_zero_filled_state(batch_size, cell.state_size, dtype)


def _generate_zero_filled_state(batch_size_tensor, state_size, dtype):
    """Generate a zero filled tensor with shape [batch_size, state_size]."""
    if batch_size_tensor is None or dtype is None:
        raise ValueError(
            'batch_size and dtype cannot be None while constructing initial state: '
            'batch_size={}, dtype={}'.format(batch_size_tensor, dtype))

    def create_zeros(unnested_state_size):
        flat_dims = tensor_shape.as_shape(unnested_state_size).as_list()
        init_state_size = [batch_size_tensor] + flat_dims
        return array_ops.zeros(init_state_size, dtype=dtype)

    if nest.is_sequence(state_size):
        return nest.map_structure(create_zeros, state_size)
    else:
        return create_zeros(state_size)
