# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Module for constructing RNN Cells."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import math

from tensorflow.contrib.compiler import jit
from tensorflow.contrib.layers.python.layers import layers
from tensorflow.contrib.rnn.python.ops import core_rnn_cell
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import op_def_registry
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.keras import activations
from tensorflow.python.keras import initializers
from tensorflow.python.keras.engine import input_spec
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import clip_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import gen_array_ops
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn_impl  # pylint: disable=unused-import
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import rnn_cell_impl
from tensorflow.python.ops import variable_scope as vs
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.util import nest


def _get_concat_variable(name, shape, dtype, num_shards):
    """Get a sharded variable concatenated into one tensor."""
    sharded_variable = _get_sharded_variable(name, shape, dtype, num_shards)
    if len(sharded_variable) == 1:
        return sharded_variable[0]

    concat_name = name + "/concat"
    concat_full_name = vs.get_variable_scope().name + "/" + concat_name + ":0"
    for value in ops.get_collection(ops.GraphKeys.CONCATENATED_VARIABLES):
        if value.name == concat_full_name:
            return value

    concat_variable = array_ops.concat(sharded_variable, 0, name=concat_name)
    ops.add_to_collection(ops.GraphKeys.CONCATENATED_VARIABLES, concat_variable)
    return concat_variable


def _get_sharded_variable(name, shape, dtype, num_shards):
    """Get a list of sharded variables with the given dtype."""
    if num_shards > shape[0]:
        raise ValueError("Too many shards: shape=%s, num_shards=%d" % (shape,
                                                                       num_shards))
    unit_shard_size = int(math.floor(shape[0] / num_shards))
    remaining_rows = shape[0] - unit_shard_size * num_shards

    shards = []
    for i in range(num_shards):
        current_size = unit_shard_size
        if i < remaining_rows:
            current_size += 1
        shards.append(
            vs.get_variable(
                name + "_%d" % i, [current_size] + shape[1:], dtype=dtype))
    return shards


def _norm(g, b, inp, scope):
    shape = inp.get_shape()[-1:]
    gamma_init = init_ops.constant_initializer(g)
    beta_init = init_ops.constant_initializer(b)
    with vs.variable_scope(scope):
        # Initialize beta and gamma for use by layer_norm.
        vs.get_variable("gamma", shape=shape, initializer=gamma_init)
        vs.get_variable("beta", shape=shape, initializer=beta_init)
    normalized = layers.layer_norm(inp, reuse=True, scope=scope)
    return normalized


class CoupledInputForgetGateLSTMCell(rnn_cell_impl.RNNCell):
    """Long short-term memory unit (LSTM) recurrent network cell.

    The default non-peephole implementation is based on:

      https://pdfs.semanticscholar.org/1154/0131eae85b2e11d53df7f1360eeb6476e7f4.pdf

    Felix Gers, Jurgen Schmidhuber, and Fred Cummins.
    "Learning to forget: Continual prediction with LSTM." IET, 850-855, 1999.

    The peephole implementation is based on:

      https://research.google.com/pubs/archive/43905.pdf

    Hasim Sak, Andrew Senior, and Francoise Beaufays.
    "Long short-term memory recurrent neural network architectures for
     large scale acoustic modeling." INTERSPEECH, 2014.

    The coupling of input and forget gate is based on:

      http://arxiv.org/pdf/1503.04069.pdf

    Greff et al. "LSTM: A Search Space Odyssey"

    The class uses optional peep-hole connections, and an optional projection
    layer.
    Layer normalization implementation is based on:

      https://arxiv.org/abs/1607.06450.

    "Layer Normalization"
    Jimmy Lei Ba, Jamie Ryan Kiros, Geoffrey E. Hinton

    and is applied before the internal nonlinearities.

    """

    def __init__(self,
                 num_units,
                 use_peepholes=False,
                 initializer=None,
                 num_proj=None,
                 proj_clip=None,
                 num_unit_shards=1,
                 num_proj_shards=1,
                 forget_bias=1.0,
                 state_is_tuple=True,
                 activation=math_ops.tanh,
                 reuse=None,
                 layer_norm=False,
                 norm_gain=1.0,
                 norm_shift=0.0):
        """Initialize the parameters for an LSTM cell.

        Args:
          num_units: int, The number of units in the LSTM cell
          use_peepholes: bool, set True to enable diagonal/peephole connections.
          initializer: (optional) The initializer to use for the weight and
            projection matrices.
          num_proj: (optional) int, The output dimensionality for the projection
            matrices.  If None, no projection is performed.
          proj_clip: (optional) A float value.  If `num_proj > 0` and `proj_clip` is
          provided, then the projected values are clipped elementwise to within
          `[-proj_clip, proj_clip]`.
          num_unit_shards: How to split the weight matrix.  If >1, the weight
            matrix is stored across num_unit_shards.
          num_proj_shards: How to split the projection matrix.  If >1, the
            projection matrix is stored across num_proj_shards.
          forget_bias: Biases of the forget gate are initialized by default to 1
            in order to reduce the scale of forgetting at the beginning of
            the training.
          state_is_tuple: If True, accepted and returned states are 2-tuples of
            the `c_state` and `m_state`.  By default (False), they are concatenated
            along the column axis.  This default behavior will soon be deprecated.
          activation: Activation function of the inner states.
          reuse: (optional) Python boolean describing whether to reuse variables
            in an existing scope.  If not `True`, and the existing scope already has
            the given variables, an error is raised.
          layer_norm: If `True`, layer normalization will be applied.
          norm_gain: float, The layer normalization gain initial value. If
            `layer_norm` has been set to `False`, this argument will be ignored.
          norm_shift: float, The layer normalization shift initial value. If
            `layer_norm` has been set to `False`, this argument will be ignored.
        """
        super(CoupledInputForgetGateLSTMCell, self).__init__(_reuse=reuse)
        if not state_is_tuple:
            logging.warn("%s: Using a concatenated state is slower and will soon be "
                         "deprecated.  Use state_is_tuple=True.", self)
        self._num_units = num_units
        self._use_peepholes = use_peepholes
        self._initializer = initializer
        self._num_proj = num_proj
        self._proj_clip = proj_clip
        self._num_unit_shards = num_unit_shards
        self._num_proj_shards = num_proj_shards
        self._forget_bias = forget_bias
        self._state_is_tuple = state_is_tuple
        self._activation = activation
        self._reuse = reuse
        self._layer_norm = layer_norm
        self._norm_gain = norm_gain
        self._norm_shift = norm_shift

        if num_proj:
            self._state_size = (
                rnn_cell_impl.LSTMStateTuple(num_units, num_proj)
                if state_is_tuple else num_units + num_proj)
            self._output_size = num_proj
        else:
            self._state_size = (
                rnn_cell_impl.LSTMStateTuple(num_units, num_units)
                if state_is_tuple else 2 * num_units)
            self._output_size = num_units

    @property
    def state_size(self):
        return self._state_size

    @property
    def output_size(self):
        return self._output_size

    def call(self, inputs, state):
        """Run one step of LSTM.

        Args:
          inputs: input Tensor, 2D, batch x num_units.
          state: if `state_is_tuple` is False, this must be a state Tensor,
            `2-D, batch x state_size`.  If `state_is_tuple` is True, this must be a
            tuple of state Tensors, both `2-D`, with column sizes `c_state` and
            `m_state`.

        Returns:
          A tuple containing:
          - A `2-D, [batch x output_dim]`, Tensor representing the output of the
            LSTM after reading `inputs` when previous state was `state`.
            Here output_dim is:
               num_proj if num_proj was set,
               num_units otherwise.
          - Tensor(s) representing the new state of LSTM after reading `inputs` when
            the previous state was `state`.  Same type and shape(s) as `state`.

        Raises:
          ValueError: If input size cannot be inferred from inputs via
            static shape inference.
        """
        sigmoid = math_ops.sigmoid

        num_proj = self._num_units if self._num_proj is None else self._num_proj

        if self._state_is_tuple:
            (c_prev, m_prev) = state
        else:
            c_prev = array_ops.slice(state, [0, 0], [-1, self._num_units])
            m_prev = array_ops.slice(state, [0, self._num_units], [-1, num_proj])

        dtype = inputs.dtype
        input_size = inputs.get_shape().with_rank(2).dims[1]
        if input_size.value is None:
            raise ValueError("Could not infer input size from inputs.get_shape()[-1]")
        concat_w = _get_concat_variable(
            "W",
            [input_size.value + num_proj, 3 * self._num_units],
            dtype,
            self._num_unit_shards)

        b = vs.get_variable(
            "B",
            shape=[3 * self._num_units],
            initializer=init_ops.zeros_initializer(),
            dtype=dtype)

        # j = new_input, f = forget_gate, o = output_gate
        cell_inputs = array_ops.concat([inputs, m_prev], 1)
        lstm_matrix = math_ops.matmul(cell_inputs, concat_w)

        # If layer nomalization is applied, do not add bias
        if not self._layer_norm:
            lstm_matrix = nn_ops.bias_add(lstm_matrix, b)

        j, f, o = array_ops.split(value=lstm_matrix, num_or_size_splits=3, axis=1)

        # Apply layer normalization
        if self._layer_norm:
            j = _norm(self._norm_gain, self._norm_shift, j, "transform")
            f = _norm(self._norm_gain, self._norm_shift, f, "forget")
            o = _norm(self._norm_gain, self._norm_shift, o, "output")

        # Diagonal connections
        if self._use_peepholes:
            w_f_diag = vs.get_variable(
                "W_F_diag", shape=[self._num_units], dtype=dtype)
            w_o_diag = vs.get_variable(
                "W_O_diag", shape=[self._num_units], dtype=dtype)

        if self._use_peepholes:
            f_act = sigmoid(f + self._forget_bias + w_f_diag * c_prev)
        else:
            f_act = sigmoid(f + self._forget_bias)
        c = (f_act * c_prev + (1 - f_act) * self._activation(j))

        # Apply layer normalization
        if self._layer_norm:
            c = _norm(self._norm_gain, self._norm_shift, c, "state")

        if self._use_peepholes:
            m = sigmoid(o + w_o_diag * c) * self._activation(c)
        else:
            m = sigmoid(o) * self._activation(c)

        if self._num_proj is not None:
            concat_w_proj = _get_concat_variable("W_P",
                                                 [self._num_units, self._num_proj],
                                                 dtype, self._num_proj_shards)

            m = math_ops.matmul(m, concat_w_proj)
            if self._proj_clip is not None:
                # pylint: disable=invalid-unary-operand-type
                m = clip_ops.clip_by_value(m, -self._proj_clip, self._proj_clip)
                # pylint: enable=invalid-unary-operand-type

        new_state = (
            rnn_cell_impl.LSTMStateTuple(c, m)
            if self._state_is_tuple else array_ops.concat([c, m], 1))
        return m, new_state


class TimeFreqLSTMCell(rnn_cell_impl.RNNCell):
    """Time-Frequency Long short-term memory unit (LSTM) recurrent network cell.

    This implementation is based on:

      Tara N. Sainath and Bo Li
      "Modeling Time-Frequency Patterns with LSTM vs. Convolutional Architectures
      for LVCSR Tasks." submitted to INTERSPEECH, 2016.

    It uses peep-hole connections and optional cell clipping.
    """

    def __init__(self,
                 num_units,
                 use_peepholes=False,
                 cell_clip=None,
                 initializer=None,
                 num_unit_shards=1,
                 forget_bias=1.0,
                 feature_size=None,
                 frequency_skip=1,
                 reuse=None):
        """Initialize the parameters for an LSTM cell.

        Args:
          num_units: int, The number of units in the LSTM cell
          use_peepholes: bool, set True to enable diagonal/peephole connections.
          cell_clip: (optional) A float value, if provided the cell state is clipped
            by this value prior to the cell output activation.
          initializer: (optional) The initializer to use for the weight and
            projection matrices.
          num_unit_shards: int, How to split the weight matrix.  If >1, the weight
            matrix is stored across num_unit_shards.
          forget_bias: float, Biases of the forget gate are initialized by default
            to 1 in order to reduce the scale of forgetting at the beginning
            of the training.
          feature_size: int, The size of the input feature the LSTM spans over.
          frequency_skip: int, The amount the LSTM filter is shifted by in
            frequency.
          reuse: (optional) Python boolean describing whether to reuse variables
            in an existing scope.  If not `True`, and the existing scope already has
            the given variables, an error is raised.
        """
        super(TimeFreqLSTMCell, self).__init__(_reuse=reuse)
        self._num_units = num_units
        self._use_peepholes = use_peepholes
        self._cell_clip = cell_clip
        self._initializer = initializer
        self._num_unit_shards = num_unit_shards
        self._forget_bias = forget_bias
        self._feature_size = feature_size
        self._frequency_skip = frequency_skip
        self._state_size = 2 * num_units
        self._output_size = num_units
        self._reuse = reuse

    @property
    def output_size(self):
        return self._output_size

    @property
    def state_size(self):
        return self._state_size

    def call(self, inputs, state):
        """Run one step of LSTM.

        Args:
          inputs: input Tensor, 2D, batch x num_units.
          state: state Tensor, 2D, batch x state_size.

        Returns:
          A tuple containing:
          - A 2D, batch x output_dim, Tensor representing the output of the LSTM
            after reading "inputs" when previous state was "state".
            Here output_dim is num_units.
          - A 2D, batch x state_size, Tensor representing the new state of LSTM
            after reading "inputs" when previous state was "state".
        Raises:
          ValueError: if an input_size was specified and the provided inputs have
            a different dimension.
        """
        sigmoid = math_ops.sigmoid
        tanh = math_ops.tanh

        freq_inputs = self._make_tf_features(inputs)
        dtype = inputs.dtype
        actual_input_size = freq_inputs[0].get_shape().as_list()[1]

        concat_w = _get_concat_variable(
            "W", [actual_input_size + 2 * self._num_units, 4 * self._num_units],
            dtype, self._num_unit_shards)

        b = vs.get_variable(
            "B",
            shape=[4 * self._num_units],
            initializer=init_ops.zeros_initializer(),
            dtype=dtype)

        # Diagonal connections
        if self._use_peepholes:
            w_f_diag = vs.get_variable(
                "W_F_diag", shape=[self._num_units], dtype=dtype)
            w_i_diag = vs.get_variable(
                "W_I_diag", shape=[self._num_units], dtype=dtype)
            w_o_diag = vs.get_variable(
                "W_O_diag", shape=[self._num_units], dtype=dtype)

        # initialize the first freq state to be zero
        m_prev_freq = array_ops.zeros(
            [inputs.shape.dims[0].value or inputs.get_shape()[0], self._num_units],
            dtype)
        for fq in range(len(freq_inputs)):
            c_prev = array_ops.slice(state, [0, 2 * fq * self._num_units],
                                     [-1, self._num_units])
            m_prev = array_ops.slice(state, [0, (2 * fq + 1) * self._num_units],
                                     [-1, self._num_units])
            # i = input_gate, j = new_input, f = forget_gate, o = output_gate
            cell_inputs = array_ops.concat([freq_inputs[fq], m_prev, m_prev_freq], 1)
            lstm_matrix = nn_ops.bias_add(math_ops.matmul(cell_inputs, concat_w), b)
            i, j, f, o = array_ops.split(
                value=lstm_matrix, num_or_size_splits=4, axis=1)

            if self._use_peepholes:
                c = (
                        sigmoid(f + self._forget_bias + w_f_diag * c_prev) * c_prev +
                        sigmoid(i + w_i_diag * c_prev) * tanh(j))
            else:
                c = (sigmoid(f + self._forget_bias) * c_prev + sigmoid(i) * tanh(j))

            if self._cell_clip is not None:
                # pylint: disable=invalid-unary-operand-type
                c = clip_ops.clip_by_value(c, -self._cell_clip, self._cell_clip)
                # pylint: enable=invalid-unary-operand-type

            if self._use_peepholes:
                m = sigmoid(o + w_o_diag * c) * tanh(c)
            else:
                m = sigmoid(o) * tanh(c)
            m_prev_freq = m
            if fq == 0:
                state_out = array_ops.concat([c, m], 1)
                m_out = m
            else:
                state_out = array_ops.concat([state_out, c, m], 1)
                m_out = array_ops.concat([m_out, m], 1)
        return m_out, state_out

    def _make_tf_features(self, input_feat):
        """Make the frequency features.

        Args:
          input_feat: input Tensor, 2D, batch x num_units.

        Returns:
          A list of frequency features, with each element containing:
          - A 2D, batch x output_dim, Tensor representing the time-frequency feature
            for that frequency index. Here output_dim is feature_size.
        Raises:
          ValueError: if input_size cannot be inferred from static shape inference.
        """
        input_size = input_feat.get_shape().with_rank(2).dims[-1].value
        if input_size is None:
            raise ValueError("Cannot infer input_size from static shape inference.")
        num_feats = int(
            (input_size - self._feature_size) / (self._frequency_skip)) + 1
        freq_inputs = []
        for f in range(num_feats):
            cur_input = array_ops.slice(input_feat, [0, f * self._frequency_skip],
                                        [-1, self._feature_size])
            freq_inputs.append(cur_input)
        return freq_inputs


class GridLSTMCell(rnn_cell_impl.RNNCell):
    """Grid Long short-term memory unit (LSTM) recurrent network cell.

    The default is based on:
      Nal Kalchbrenner, Ivo Danihelka and Alex Graves
      "Grid Long Short-Term Memory," Proc. ICLR 2016.
      http://arxiv.org/abs/1507.01526

    When peephole connections are used, the implementation is based on:
      Tara N. Sainath and Bo Li
      "Modeling Time-Frequency Patterns with LSTM vs. Convolutional Architectures
      for LVCSR Tasks." submitted to INTERSPEECH, 2016.

    The code uses optional peephole connections, shared_weights and cell clipping.
    """

    def __init__(self,
                 num_units,
                 use_peepholes=False,
                 share_time_frequency_weights=False,
                 cell_clip=None,
                 initializer=None,
                 num_unit_shards=1,
                 forget_bias=1.0,
                 feature_size=None,
                 frequency_skip=None,
                 num_frequency_blocks=None,
                 start_freqindex_list=None,
                 end_freqindex_list=None,
                 couple_input_forget_gates=False,
                 state_is_tuple=True,
                 reuse=None):
        """Initialize the parameters for an LSTM cell.

        Args:
          num_units: int, The number of units in the LSTM cell
          use_peepholes: (optional) bool, default False. Set True to enable
            diagonal/peephole connections.
          share_time_frequency_weights: (optional) bool, default False. Set True to
            enable shared cell weights between time and frequency LSTMs.
          cell_clip: (optional) A float value, default None, if provided the cell
            state is clipped by this value prior to the cell output activation.
          initializer: (optional) The initializer to use for the weight and
            projection matrices, default None.
          num_unit_shards: (optional) int, default 1, How to split the weight
            matrix. If > 1, the weight matrix is stored across num_unit_shards.
          forget_bias: (optional) float, default 1.0, The initial bias of the
            forget gates, used to reduce the scale of forgetting at the beginning
            of the training.
          feature_size: (optional) int, default None, The size of the input feature
            the LSTM spans over.
          frequency_skip: (optional) int, default None, The amount the LSTM filter
            is shifted by in frequency.
          num_frequency_blocks: [required] A list of frequency blocks needed to
            cover the whole input feature splitting defined by start_freqindex_list
            and end_freqindex_list.
          start_freqindex_list: [optional], list of ints, default None,  The
            starting frequency index for each frequency block.
          end_freqindex_list: [optional], list of ints, default None. The ending
            frequency index for each frequency block.
          couple_input_forget_gates: (optional) bool, default False, Whether to
            couple the input and forget gates, i.e. f_gate = 1.0 - i_gate, to reduce
            model parameters and computation cost.
          state_is_tuple: If True, accepted and returned states are 2-tuples of
            the `c_state` and `m_state`.  By default (False), they are concatenated
            along the column axis.  This default behavior will soon be deprecated.
          reuse: (optional) Python boolean describing whether to reuse variables
            in an existing scope.  If not `True`, and the existing scope already has
            the given variables, an error is raised.
        Raises:
          ValueError: if the num_frequency_blocks list is not specified
        """
        super(GridLSTMCell, self).__init__(_reuse=reuse)
        if not state_is_tuple:
            logging.warn("%s: Using a concatenated state is slower and will soon be "
                         "deprecated.  Use state_is_tuple=True.", self)
        self._num_units = num_units
        self._use_peepholes = use_peepholes
        self._share_time_frequency_weights = share_time_frequency_weights
        self._couple_input_forget_gates = couple_input_forget_gates
        self._state_is_tuple = state_is_tuple
        self._cell_clip = cell_clip
        self._initializer = initializer
        self._num_unit_shards = num_unit_shards
        self._forget_bias = forget_bias
        self._feature_size = feature_size
        self._frequency_skip = frequency_skip
        self._start_freqindex_list = start_freqindex_list
        self._end_freqindex_list = end_freqindex_list
        self._num_frequency_blocks = num_frequency_blocks
        self._total_blocks = 0
        self._reuse = reuse
        if self._num_frequency_blocks is None:
            raise ValueError("Must specify num_frequency_blocks")

        for block_index in range(len(self._num_frequency_blocks)):
            self._total_blocks += int(self._num_frequency_blocks[block_index])
        if state_is_tuple:
            state_names = ""
            for block_index in range(len(self._num_frequency_blocks)):
                for freq_index in range(self._num_frequency_blocks[block_index]):
                    name_prefix = "state_f%02d_b%02d" % (freq_index, block_index)
                    state_names += ("%s_c, %s_m," % (name_prefix, name_prefix))
            self._state_tuple_type = collections.namedtuple("GridLSTMStateTuple",
                                                            state_names.strip(","))
            self._state_size = self._state_tuple_type(*(
                    [num_units, num_units] * self._total_blocks))
        else:
            self._state_tuple_type = None
            self._state_size = num_units * self._total_blocks * 2
        self._output_size = num_units * self._total_blocks * 2

    @property
    def output_size(self):
        return self._output_size

    @property
    def state_size(self):
        return self._state_size

    @property
    def state_tuple_type(self):
        return self._state_tuple_type

    def call(self, inputs, state):
        """Run one step of LSTM.

        Args:
          inputs: input Tensor, 2D, [batch, feature_size].
          state: Tensor or tuple of Tensors, 2D, [batch, state_size], depends on the
            flag self._state_is_tuple.

        Returns:
          A tuple containing:
          - A 2D, [batch, output_dim], Tensor representing the output of the LSTM
            after reading "inputs" when previous state was "state".
            Here output_dim is num_units.
          - A 2D, [batch, state_size], Tensor representing the new state of LSTM
            after reading "inputs" when previous state was "state".
        Raises:
          ValueError: if an input_size was specified and the provided inputs have
            a different dimension.
        """
        batch_size = tensor_shape.dimension_value(
            inputs.shape[0]) or array_ops.shape(inputs)[0]
        freq_inputs = self._make_tf_features(inputs)
        m_out_lst = []
        state_out_lst = []
        for block in range(len(freq_inputs)):
            m_out_lst_current, state_out_lst_current = self._compute(
                freq_inputs[block],
                block,
                state,
                batch_size,
                state_is_tuple=self._state_is_tuple)
            m_out_lst.extend(m_out_lst_current)
            state_out_lst.extend(state_out_lst_current)
        if self._state_is_tuple:
            state_out = self._state_tuple_type(*state_out_lst)
        else:
            state_out = array_ops.concat(state_out_lst, 1)
        m_out = array_ops.concat(m_out_lst, 1)
        return m_out, state_out

    def _compute(self,
                 freq_inputs,
                 block,
                 state,
                 batch_size,
                 state_prefix="state",
                 state_is_tuple=True):
        """Run the actual computation of one step LSTM.

        Args:
          freq_inputs: list of Tensors, 2D, [batch, feature_size].
          block: int, current frequency block index to process.
          state: Tensor or tuple of Tensors, 2D, [batch, state_size], it depends on
            the flag state_is_tuple.
          batch_size: int32, batch size.
          state_prefix: (optional) string, name prefix for states, defaults to
            "state".
          state_is_tuple: boolean, indicates whether the state is a tuple or Tensor.

        Returns:
          A tuple, containing:
          - A list of [batch, output_dim] Tensors, representing the output of the
            LSTM given the inputs and state.
          - A list of [batch, state_size] Tensors, representing the LSTM state
            values given the inputs and previous state.
        """
        sigmoid = math_ops.sigmoid
        tanh = math_ops.tanh
        num_gates = 3 if self._couple_input_forget_gates else 4
        dtype = freq_inputs[0].dtype
        actual_input_size = freq_inputs[0].get_shape().as_list()[1]

        concat_w_f = _get_concat_variable(
            "W_f_%d" % block,
            [actual_input_size + 2 * self._num_units, num_gates * self._num_units],
            dtype, self._num_unit_shards)
        b_f = vs.get_variable(
            "B_f_%d" % block,
            shape=[num_gates * self._num_units],
            initializer=init_ops.zeros_initializer(),
            dtype=dtype)
        if not self._share_time_frequency_weights:
            concat_w_t = _get_concat_variable("W_t_%d" % block, [
                actual_input_size + 2 * self._num_units, num_gates * self._num_units
            ], dtype, self._num_unit_shards)
            b_t = vs.get_variable(
                "B_t_%d" % block,
                shape=[num_gates * self._num_units],
                initializer=init_ops.zeros_initializer(),
                dtype=dtype)

        if self._use_peepholes:
            # Diagonal connections
            if not self._couple_input_forget_gates:
                w_f_diag_freqf = vs.get_variable(
                    "W_F_diag_freqf_%d" % block, shape=[self._num_units], dtype=dtype)
                w_f_diag_freqt = vs.get_variable(
                    "W_F_diag_freqt_%d" % block, shape=[self._num_units], dtype=dtype)
            w_i_diag_freqf = vs.get_variable(
                "W_I_diag_freqf_%d" % block, shape=[self._num_units], dtype=dtype)
            w_i_diag_freqt = vs.get_variable(
                "W_I_diag_freqt_%d" % block, shape=[self._num_units], dtype=dtype)
            w_o_diag_freqf = vs.get_variable(
                "W_O_diag_freqf_%d" % block, shape=[self._num_units], dtype=dtype)
            w_o_diag_freqt = vs.get_variable(
                "W_O_diag_freqt_%d" % block, shape=[self._num_units], dtype=dtype)
            if not self._share_time_frequency_weights:
                if not self._couple_input_forget_gates:
                    w_f_diag_timef = vs.get_variable(
                        "W_F_diag_timef_%d" % block, shape=[self._num_units], dtype=dtype)
                    w_f_diag_timet = vs.get_variable(
                        "W_F_diag_timet_%d" % block, shape=[self._num_units], dtype=dtype)
                w_i_diag_timef = vs.get_variable(
                    "W_I_diag_timef_%d" % block, shape=[self._num_units], dtype=dtype)
                w_i_diag_timet = vs.get_variable(
                    "W_I_diag_timet_%d" % block, shape=[self._num_units], dtype=dtype)
                w_o_diag_timef = vs.get_variable(
                    "W_O_diag_timef_%d" % block, shape=[self._num_units], dtype=dtype)
                w_o_diag_timet = vs.get_variable(
                    "W_O_diag_timet_%d" % block, shape=[self._num_units], dtype=dtype)

        # initialize the first freq state to be zero
        m_prev_freq = array_ops.zeros([batch_size, self._num_units], dtype)
        c_prev_freq = array_ops.zeros([batch_size, self._num_units], dtype)
        for freq_index in range(len(freq_inputs)):
            if state_is_tuple:
                name_prefix = "%s_f%02d_b%02d" % (state_prefix, freq_index, block)
                c_prev_time = getattr(state, name_prefix + "_c")
                m_prev_time = getattr(state, name_prefix + "_m")
            else:
                c_prev_time = array_ops.slice(
                    state, [0, 2 * freq_index * self._num_units], [-1, self._num_units])
                m_prev_time = array_ops.slice(
                    state, [0, (2 * freq_index + 1) * self._num_units],
                    [-1, self._num_units])

            # i = input_gate, j = new_input, f = forget_gate, o = output_gate
            cell_inputs = array_ops.concat(
                [freq_inputs[freq_index], m_prev_time, m_prev_freq], 1)

            # F-LSTM
            lstm_matrix_freq = nn_ops.bias_add(
                math_ops.matmul(cell_inputs, concat_w_f), b_f)
            if self._couple_input_forget_gates:
                i_freq, j_freq, o_freq = array_ops.split(
                    value=lstm_matrix_freq, num_or_size_splits=num_gates, axis=1)
                f_freq = None
            else:
                i_freq, j_freq, f_freq, o_freq = array_ops.split(
                    value=lstm_matrix_freq, num_or_size_splits=num_gates, axis=1)
            # T-LSTM
            if self._share_time_frequency_weights:
                i_time = i_freq
                j_time = j_freq
                f_time = f_freq
                o_time = o_freq
            else:
                lstm_matrix_time = nn_ops.bias_add(
                    math_ops.matmul(cell_inputs, concat_w_t), b_t)
                if self._couple_input_forget_gates:
                    i_time, j_time, o_time = array_ops.split(
                        value=lstm_matrix_time, num_or_size_splits=num_gates, axis=1)
                    f_time = None
                else:
                    i_time, j_time, f_time, o_time = array_ops.split(
                        value=lstm_matrix_time, num_or_size_splits=num_gates, axis=1)

            # F-LSTM c_freq
            # input gate activations
            if self._use_peepholes:
                i_freq_g = sigmoid(i_freq + w_i_diag_freqf * c_prev_freq +
                                   w_i_diag_freqt * c_prev_time)
            else:
                i_freq_g = sigmoid(i_freq)
            # forget gate activations
            if self._couple_input_forget_gates:
                f_freq_g = 1.0 - i_freq_g
            else:
                if self._use_peepholes:
                    f_freq_g = sigmoid(f_freq + self._forget_bias + w_f_diag_freqf *
                                       c_prev_freq + w_f_diag_freqt * c_prev_time)
                else:
                    f_freq_g = sigmoid(f_freq + self._forget_bias)
            # cell state
            c_freq = f_freq_g * c_prev_freq + i_freq_g * tanh(j_freq)
            if self._cell_clip is not None:
                # pylint: disable=invalid-unary-operand-type
                c_freq = clip_ops.clip_by_value(c_freq, -self._cell_clip,
                                                self._cell_clip)
                # pylint: enable=invalid-unary-operand-type

            # T-LSTM c_freq
            # input gate activations
            if self._use_peepholes:
                if self._share_time_frequency_weights:
                    i_time_g = sigmoid(i_time + w_i_diag_freqf * c_prev_freq +
                                       w_i_diag_freqt * c_prev_time)
                else:
                    i_time_g = sigmoid(i_time + w_i_diag_timef * c_prev_freq +
                                       w_i_diag_timet * c_prev_time)
            else:
                i_time_g = sigmoid(i_time)
            # forget gate activations
            if self._couple_input_forget_gates:
                f_time_g = 1.0 - i_time_g
            else:
                if self._use_peepholes:
                    if self._share_time_frequency_weights:
                        f_time_g = sigmoid(f_time + self._forget_bias + w_f_diag_freqf *
                                           c_prev_freq + w_f_diag_freqt * c_prev_time)
                    else:
                        f_time_g = sigmoid(f_time + self._forget_bias + w_f_diag_timef *
                                           c_prev_freq + w_f_diag_timet * c_prev_time)
                else:
                    f_time_g = sigmoid(f_time + self._forget_bias)
            # cell state
            c_time = f_time_g * c_prev_time + i_time_g * tanh(j_time)
            if self._cell_clip is not None:
                # pylint: disable=invalid-unary-operand-type
                c_time = clip_ops.clip_by_value(c_time, -self._cell_clip,
                                                self._cell_clip)
                # pylint: enable=invalid-unary-operand-type

            # F-LSTM m_freq
            if self._use_peepholes:
                m_freq = sigmoid(o_freq + w_o_diag_freqf * c_freq +
                                 w_o_diag_freqt * c_time) * tanh(c_freq)
            else:
                m_freq = sigmoid(o_freq) * tanh(c_freq)

            # T-LSTM m_time
            if self._use_peepholes:
                if self._share_time_frequency_weights:
                    m_time = sigmoid(o_time + w_o_diag_freqf * c_freq +
                                     w_o_diag_freqt * c_time) * tanh(c_time)
                else:
                    m_time = sigmoid(o_time + w_o_diag_timef * c_freq +
                                     w_o_diag_timet * c_time) * tanh(c_time)
            else:
                m_time = sigmoid(o_time) * tanh(c_time)

            m_prev_freq = m_freq
            c_prev_freq = c_freq
            # Concatenate the outputs for T-LSTM and F-LSTM for each shift
            if freq_index == 0:
                state_out_lst = [c_time, m_time]
                m_out_lst = [m_time, m_freq]
            else:
                state_out_lst.extend([c_time, m_time])
                m_out_lst.extend([m_time, m_freq])

        return m_out_lst, state_out_lst

    def _make_tf_features(self, input_feat, slice_offset=0):
        """Make the frequency features.

        Args:
          input_feat: input Tensor, 2D, [batch, num_units].
          slice_offset: (optional) Python int, default 0, the slicing offset is only
            used for the backward processing in the BidirectionalGridLSTMCell. It
            specifies a different starting point instead of always 0 to enable the
            forward and backward processing look at different frequency blocks.

        Returns:
          A list of frequency features, with each element containing:
          - A 2D, [batch, output_dim], Tensor representing the time-frequency
            feature for that frequency index. Here output_dim is feature_size.
        Raises:
          ValueError: if input_size cannot be inferred from static shape inference.
        """
        input_size = input_feat.get_shape().with_rank(2).dims[-1].value
        if input_size is None:
            raise ValueError("Cannot infer input_size from static shape inference.")
        if slice_offset > 0:
            # Padding to the end
            inputs = array_ops.pad(input_feat,
                                   array_ops.constant(
                                       [0, 0, 0, slice_offset],
                                       shape=[2, 2],
                                       dtype=dtypes.int32), "CONSTANT")
        elif slice_offset < 0:
            # Padding to the front
            inputs = array_ops.pad(input_feat,
                                   array_ops.constant(
                                       [0, 0, -slice_offset, 0],
                                       shape=[2, 2],
                                       dtype=dtypes.int32), "CONSTANT")
            slice_offset = 0
        else:
            inputs = input_feat
        freq_inputs = []
        if not self._start_freqindex_list:
            if len(self._num_frequency_blocks) != 1:
                raise ValueError("Length of num_frequency_blocks"
                                 " is not 1, but instead is %d" %
                                 len(self._num_frequency_blocks))
            num_feats = int(
                (input_size - self._feature_size) / (self._frequency_skip)) + 1
            if num_feats != self._num_frequency_blocks[0]:
                raise ValueError(
                    "Invalid num_frequency_blocks, requires %d but gets %d, please"
                    " check the input size and filter config are correct." %
                    (self._num_frequency_blocks[0], num_feats))
            block_inputs = []
            for f in range(num_feats):
                cur_input = array_ops.slice(
                    inputs, [0, slice_offset + f * self._frequency_skip],
                    [-1, self._feature_size])
                block_inputs.append(cur_input)
            freq_inputs.append(block_inputs)
        else:
            if len(self._start_freqindex_list) != len(self._end_freqindex_list):
                raise ValueError("Length of start and end freqindex_list"
                                 " does not match %d %d",
                                 len(self._start_freqindex_list),
                                 len(self._end_freqindex_list))
            if len(self._num_frequency_blocks) != len(self._start_freqindex_list):
                raise ValueError("Length of num_frequency_blocks"
                                 " is not equal to start_freqindex_list %d %d",
                                 len(self._num_frequency_blocks),
                                 len(self._start_freqindex_list))
            for b in range(len(self._start_freqindex_list)):
                start_index = self._start_freqindex_list[b]
                end_index = self._end_freqindex_list[b]
                cur_size = end_index - start_index
                block_feats = int(
                    (cur_size - self._feature_size) / (self._frequency_skip)) + 1
                if block_feats != self._num_frequency_blocks[b]:
                    raise ValueError(
                        "Invalid num_frequency_blocks, requires %d but gets %d, please"
                        " check the input size and filter config are correct." %
                        (self._num_frequency_blocks[b], block_feats))
                block_inputs = []
                for f in range(block_feats):
                    cur_input = array_ops.slice(
                        inputs,
                        [0, start_index + slice_offset + f * self._frequency_skip],
                        [-1, self._feature_size])
                    block_inputs.append(cur_input)
                freq_inputs.append(block_inputs)
        return freq_inputs


class BidirectionalGridLSTMCell(GridLSTMCell):
    """Bidirectional GridLstm cell.

    The bidirection connection is only used in the frequency direction, which
    hence doesn't affect the time direction's real-time processing that is
    required for online recognition systems.
    The current implementation uses different weights for the two directions.
    """

    def __init__(self,
                 num_units,
                 use_peepholes=False,
                 share_time_frequency_weights=False,
                 cell_clip=None,
                 initializer=None,
                 num_unit_shards=1,
                 forget_bias=1.0,
                 feature_size=None,
                 frequency_skip=None,
                 num_frequency_blocks=None,
                 start_freqindex_list=None,
                 end_freqindex_list=None,
                 couple_input_forget_gates=False,
                 backward_slice_offset=0,
                 reuse=None):
        """Initialize the parameters for an LSTM cell.

        Args:
          num_units: int, The number of units in the LSTM cell
          use_peepholes: (optional) bool, default False. Set True to enable
            diagonal/peephole connections.
          share_time_frequency_weights: (optional) bool, default False. Set True to
            enable shared cell weights between time and frequency LSTMs.
          cell_clip: (optional) A float value, default None, if provided the cell
            state is clipped by this value prior to the cell output activation.
          initializer: (optional) The initializer to use for the weight and
            projection matrices, default None.
          num_unit_shards: (optional) int, default 1, How to split the weight
            matrix. If > 1, the weight matrix is stored across num_unit_shards.
          forget_bias: (optional) float, default 1.0, The initial bias of the
            forget gates, used to reduce the scale of forgetting at the beginning
            of the training.
          feature_size: (optional) int, default None, The size of the input feature
            the LSTM spans over.
          frequency_skip: (optional) int, default None, The amount the LSTM filter
            is shifted by in frequency.
          num_frequency_blocks: [required] A list of frequency blocks needed to
            cover the whole input feature splitting defined by start_freqindex_list
            and end_freqindex_list.
          start_freqindex_list: [optional], list of ints, default None,  The
            starting frequency index for each frequency block.
          end_freqindex_list: [optional], list of ints, default None. The ending
            frequency index for each frequency block.
          couple_input_forget_gates: (optional) bool, default False, Whether to
            couple the input and forget gates, i.e. f_gate = 1.0 - i_gate, to reduce
            model parameters and computation cost.
          backward_slice_offset: (optional) int32, default 0, the starting offset to
            slice the feature for backward processing.
          reuse: (optional) Python boolean describing whether to reuse variables
            in an existing scope.  If not `True`, and the existing scope already has
            the given variables, an error is raised.
        """
        super(BidirectionalGridLSTMCell, self).__init__(
            num_units, use_peepholes, share_time_frequency_weights, cell_clip,
            initializer, num_unit_shards, forget_bias, feature_size, frequency_skip,
            num_frequency_blocks, start_freqindex_list, end_freqindex_list,
            couple_input_forget_gates, True, reuse)
        self._backward_slice_offset = int(backward_slice_offset)
        state_names = ""
        for direction in ["fwd", "bwd"]:
            for block_index in range(len(self._num_frequency_blocks)):
                for freq_index in range(self._num_frequency_blocks[block_index]):
                    name_prefix = "%s_state_f%02d_b%02d" % (direction, freq_index,
                                                            block_index)
                    state_names += ("%s_c, %s_m," % (name_prefix, name_prefix))
        self._state_tuple_type = collections.namedtuple(
            "BidirectionalGridLSTMStateTuple", state_names.strip(","))
        self._state_size = self._state_tuple_type(*(
                [num_units, num_units] * self._total_blocks * 2))
        self._output_size = 2 * num_units * self._total_blocks * 2

    def call(self, inputs, state):
        """Run one step of LSTM.

        Args:
          inputs: input Tensor, 2D, [batch, num_units].
          state: tuple of Tensors, 2D, [batch, state_size].

        Returns:
          A tuple containing:
          - A 2D, [batch, output_dim], Tensor representing the output of the LSTM
            after reading "inputs" when previous state was "state".
            Here output_dim is num_units.
          - A 2D, [batch, state_size], Tensor representing the new state of LSTM
            after reading "inputs" when previous state was "state".
        Raises:
          ValueError: if an input_size was specified and the provided inputs have
            a different dimension.
        """
        batch_size = tensor_shape.dimension_value(
            inputs.shape[0]) or array_ops.shape(inputs)[0]
        fwd_inputs = self._make_tf_features(inputs)
        if self._backward_slice_offset:
            bwd_inputs = self._make_tf_features(inputs, self._backward_slice_offset)
        else:
            bwd_inputs = fwd_inputs

        # Forward processing
        with vs.variable_scope("fwd"):
            fwd_m_out_lst = []
            fwd_state_out_lst = []
            for block in range(len(fwd_inputs)):
                fwd_m_out_lst_current, fwd_state_out_lst_current = self._compute(
                    fwd_inputs[block],
                    block,
                    state,
                    batch_size,
                    state_prefix="fwd_state",
                    state_is_tuple=True)
                fwd_m_out_lst.extend(fwd_m_out_lst_current)
                fwd_state_out_lst.extend(fwd_state_out_lst_current)
        # Backward processing
        bwd_m_out_lst = []
        bwd_state_out_lst = []
        with vs.variable_scope("bwd"):
            for block in range(len(bwd_inputs)):
                # Reverse the blocks
                bwd_inputs_reverse = bwd_inputs[block][::-1]
                bwd_m_out_lst_current, bwd_state_out_lst_current = self._compute(
                    bwd_inputs_reverse,
                    block,
                    state,
                    batch_size,
                    state_prefix="bwd_state",
                    state_is_tuple=True)
                bwd_m_out_lst.extend(bwd_m_out_lst_current)
                bwd_state_out_lst.extend(bwd_state_out_lst_current)
        state_out = self._state_tuple_type(*(fwd_state_out_lst + bwd_state_out_lst))
        # Outputs are always concated as it is never used separately.
        m_out = array_ops.concat(fwd_m_out_lst + bwd_m_out_lst, 1)
        return m_out, state_out


# pylint: disable=protected-access
_Linear = core_rnn_cell._Linear  # pylint: disable=invalid-name


# pylint: enable=protected-access


class AttentionCellWrapper(rnn_cell_impl.RNNCell):
    """Basic attention cell wrapper.

    Implementation based on https://arxiv.org/abs/1601.06733.
    """

    def __init__(self,
                 cell,
                 attn_length,
                 attn_size=None,
                 attn_vec_size=None,
                 input_size=None,
                 state_is_tuple=True,
                 reuse=None):
        """Create a cell with attention.

        Args:
          cell: an RNNCell, an attention is added to it.
          attn_length: integer, the size of an attention window.
          attn_size: integer, the size of an attention vector. Equal to
              cell.output_size by default.
          attn_vec_size: integer, the number of convolutional features calculated
              on attention state and a size of the hidden layer built from
              base cell state. Equal attn_size to by default.
          input_size: integer, the size of a hidden linear layer,
              built from inputs and attention. Derived from the input tensor
              by default.
          state_is_tuple: If True, accepted and returned states are n-tuples, where
            `n = len(cells)`.  By default (False), the states are all
            concatenated along the column axis.
          reuse: (optional) Python boolean describing whether to reuse variables
            in an existing scope.  If not `True`, and the existing scope already has
            the given variables, an error is raised.

        Raises:
          TypeError: if cell is not an RNNCell.
          ValueError: if cell returns a state tuple but the flag
              `state_is_tuple` is `False` or if attn_length is zero or less.
        """
        super(AttentionCellWrapper, self).__init__(_reuse=reuse)
        rnn_cell_impl.assert_like_rnncell("cell", cell)
        if nest.is_sequence(cell.state_size) and not state_is_tuple:
            raise ValueError(
                "Cell returns tuple of states, but the flag "
                "state_is_tuple is not set. State size is: %s" % str(cell.state_size))
        if attn_length <= 0:
            raise ValueError(
                "attn_length should be greater than zero, got %s" % str(attn_length))
        if not state_is_tuple:
            logging.warn("%s: Using a concatenated state is slower and will soon be "
                         "deprecated.  Use state_is_tuple=True.", self)
        if attn_size is None:
            attn_size = cell.output_size
        if attn_vec_size is None:
            attn_vec_size = attn_size
        self._state_is_tuple = state_is_tuple
        self._cell = cell
        self._attn_vec_size = attn_vec_size
        self._input_size = input_size
        self._attn_size = attn_size
        self._attn_length = attn_length
        self._reuse = reuse
        self._linear1 = None
        self._linear2 = None
        self._linear3 = None

    @property
    def state_size(self):
        size = (self._cell.state_size, self._attn_size,
                self._attn_size * self._attn_length)
        if self._state_is_tuple:
            return size
        else:
            return sum(list(size))

    @property
    def output_size(self):
        return self._attn_size

    def call(self, inputs, state):
        """Long short-term memory cell with attention (LSTMA)."""
        if self._state_is_tuple:
            state, attns, attn_states = state
        else:
            states = state
            state = array_ops.slice(states, [0, 0], [-1, self._cell.state_size])
            attns = array_ops.slice(states, [0, self._cell.state_size],
                                    [-1, self._attn_size])
            attn_states = array_ops.slice(
                states, [0, self._cell.state_size + self._attn_size],
                [-1, self._attn_size * self._attn_length])
        attn_states = array_ops.reshape(attn_states,
                                        [-1, self._attn_length, self._attn_size])
        input_size = self._input_size
        if input_size is None:
            input_size = inputs.get_shape().as_list()[1]
        if self._linear1 is None:
            self._linear1 = _Linear([inputs, attns], input_size, True)
        inputs = self._linear1([inputs, attns])
        cell_output, new_state = self._cell(inputs, state)
        if self._state_is_tuple:
            new_state_cat = array_ops.concat(nest.flatten(new_state), 1)
        else:
            new_state_cat = new_state
        new_attns, new_attn_states = self._attention(new_state_cat, attn_states)
        with vs.variable_scope("attn_output_projection"):
            if self._linear2 is None:
                self._linear2 = _Linear([cell_output, new_attns], self._attn_size, True)
            output = self._linear2([cell_output, new_attns])
        new_attn_states = array_ops.concat(
            [new_attn_states, array_ops.expand_dims(output, 1)], 1)
        new_attn_states = array_ops.reshape(
            new_attn_states, [-1, self._attn_length * self._attn_size])
        new_state = (new_state, new_attns, new_attn_states)
        if not self._state_is_tuple:
            new_state = array_ops.concat(list(new_state), 1)
        return output, new_state

    def _attention(self, query, attn_states):
        conv2d = nn_ops.conv2d
        reduce_sum = math_ops.reduce_sum
        softmax = nn_ops.softmax
        tanh = math_ops.tanh

        with vs.variable_scope("attention"):
            k = vs.get_variable("attn_w",
                                [1, 1, self._attn_size, self._attn_vec_size])
            v = vs.get_variable("attn_v", [self._attn_vec_size])
            hidden = array_ops.reshape(attn_states,
                                       [-1, self._attn_length, 1, self._attn_size])
            hidden_features = conv2d(hidden, k, [1, 1, 1, 1], "SAME")
            if self._linear3 is None:
                self._linear3 = _Linear(query, self._attn_vec_size, True)
            y = self._linear3(query)
            y = array_ops.reshape(y, [-1, 1, 1, self._attn_vec_size])
            s = reduce_sum(v * tanh(hidden_features + y), [2, 3])
            a = softmax(s)
            d = reduce_sum(
                array_ops.reshape(a, [-1, self._attn_length, 1, 1]) * hidden, [1, 2])
            new_attns = array_ops.reshape(d, [-1, self._attn_size])
            new_attn_states = array_ops.slice(attn_states, [0, 1, 0], [-1, -1, -1])
            return new_attns, new_attn_states


class HighwayWrapper(rnn_cell_impl.RNNCell):
    """RNNCell wrapper that adds highway connection on cell input and output.

    Based on:
      R. K. Srivastava, K. Greff, and J. Schmidhuber, "Highway networks",
      arXiv preprint arXiv:1505.00387, 2015.
      https://arxiv.org/abs/1505.00387
    """

    def __init__(self,
                 cell,
                 couple_carry_transform_gates=True,
                 carry_bias_init=1.0):
        """Constructs a `HighwayWrapper` for `cell`.

        Args:
          cell: An instance of `RNNCell`.
          couple_carry_transform_gates: boolean, should the Carry and Transform gate
            be coupled.
          carry_bias_init: float, carry gates bias initialization.
        """
        self._cell = cell
        self._couple_carry_transform_gates = couple_carry_transform_gates
        self._carry_bias_init = carry_bias_init

    @property
    def state_size(self):
        return self._cell.state_size

    @property
    def output_size(self):
        return self._cell.output_size

    def zero_state(self, batch_size, dtype):
        with ops.name_scope(type(self).__name__ + "ZeroState", values=[batch_size]):
            return self._cell.zero_state(batch_size, dtype)

    def _highway(self, inp, out):
        input_size = inp.get_shape().with_rank(2).dims[1].value
        carry_weight = vs.get_variable("carry_w", [input_size, input_size])
        carry_bias = vs.get_variable(
            "carry_b", [input_size],
            initializer=init_ops.constant_initializer(self._carry_bias_init))
        carry = math_ops.sigmoid(nn_ops.xw_plus_b(inp, carry_weight, carry_bias))
        if self._couple_carry_transform_gates:
            transform = 1 - carry
        else:
            transform_weight = vs.get_variable("transform_w",
                                               [input_size, input_size])
            transform_bias = vs.get_variable(
                "transform_b", [input_size],
                initializer=init_ops.constant_initializer(-self._carry_bias_init))
            transform = math_ops.sigmoid(
                nn_ops.xw_plus_b(inp, transform_weight, transform_bias))
        return inp * carry + out * transform

    def __call__(self, inputs, state, scope=None):
        """Run the cell and add its inputs to its outputs.

        Args:
          inputs: cell inputs.
          state: cell state.
          scope: optional cell scope.

        Returns:
          Tuple of cell outputs and new state.

        Raises:
          TypeError: If cell inputs and outputs have different structure (type).
          ValueError: If cell inputs and outputs have different structure (value).
        """
        outputs, new_state = self._cell(inputs, state, scope=scope)
        nest.assert_same_structure(inputs, outputs)

        # Ensure shapes match
        def assert_shape_match(inp, out):
            inp.get_shape().assert_is_compatible_with(out.get_shape())

        nest.map_structure(assert_shape_match, inputs, outputs)
        res_outputs = nest.map_structure(self._highway, inputs, outputs)
        return (res_outputs, new_state)


class LayerNormBasicLSTMCell(rnn_cell_impl.RNNCell):
    """LSTM unit with layer normalization and recurrent dropout.

    This class adds layer normalization and recurrent dropout to a
    basic LSTM unit. Layer normalization implementation is based on:

      https://arxiv.org/abs/1607.06450.

    "Layer Normalization"
    Jimmy Lei Ba, Jamie Ryan Kiros, Geoffrey E. Hinton

    and is applied before the internal nonlinearities.
    Recurrent dropout is base on:

      https://arxiv.org/abs/1603.05118

    "Recurrent Dropout without Memory Loss"
    Stanislau Semeniuta, Aliaksei Severyn, Erhardt Barth.
    """

    def __init__(self,
                 num_units,
                 forget_bias=1.0,
                 input_size=None,
                 activation=math_ops.tanh,
                 layer_norm=True,
                 norm_gain=1.0,
                 norm_shift=0.0,
                 dropout_keep_prob=1.0,
                 dropout_prob_seed=None,
                 reuse=None):
        """Initializes the basic LSTM cell.

        Args:
          num_units: int, The number of units in the LSTM cell.
          forget_bias: float, The bias added to forget gates (see above).
          input_size: Deprecated and unused.
          activation: Activation function of the inner states.
          layer_norm: If `True`, layer normalization will be applied.
          norm_gain: float, The layer normalization gain initial value. If
            `layer_norm` has been set to `False`, this argument will be ignored.
          norm_shift: float, The layer normalization shift initial value. If
            `layer_norm` has been set to `False`, this argument will be ignored.
          dropout_keep_prob: unit Tensor or float between 0 and 1 representing the
            recurrent dropout probability value. If float and 1.0, no dropout will
            be applied.
          dropout_prob_seed: (optional) integer, the randomness seed.
          reuse: (optional) Python boolean describing whether to reuse variables
            in an existing scope.  If not `True`, and the existing scope already has
            the given variables, an error is raised.
        """
        super(LayerNormBasicLSTMCell, self).__init__(_reuse=reuse)

        if input_size is not None:
            logging.warn("%s: The input_size parameter is deprecated.", self)

        self._num_units = num_units
        self._activation = activation
        self._forget_bias = forget_bias
        self._keep_prob = dropout_keep_prob
        self._seed = dropout_prob_seed
        self._layer_norm = layer_norm
        self._norm_gain = norm_gain
        self._norm_shift = norm_shift
        self._reuse = reuse

    @property
    def state_size(self):
        return rnn_cell_impl.LSTMStateTuple(self._num_units, self._num_units)

    @property
    def output_size(self):
        return self._num_units

    def _norm(self, inp, scope, dtype=dtypes.float32):
        shape = inp.get_shape()[-1:]
        gamma_init = init_ops.constant_initializer(self._norm_gain)
        beta_init = init_ops.constant_initializer(self._norm_shift)
        with vs.variable_scope(scope):
            # Initialize beta and gamma for use by layer_norm.
            vs.get_variable("gamma", shape=shape, initializer=gamma_init, dtype=dtype)
            vs.get_variable("beta", shape=shape, initializer=beta_init, dtype=dtype)
        normalized = layers.layer_norm(inp, reuse=True, scope=scope)
        return normalized

    def _linear(self, args):
        out_size = 4 * self._num_units
        proj_size = args.get_shape()[-1]
        dtype = args.dtype
        weights = vs.get_variable("kernel", [proj_size, out_size], dtype=dtype)
        out = math_ops.matmul(args, weights)
        if not self._layer_norm:
            bias = vs.get_variable("bias", [out_size], dtype=dtype)
            out = nn_ops.bias_add(out, bias)
        return out

    def call(self, inputs, state):
        """LSTM cell with layer normalization and recurrent dropout."""
        c, h = state
        args = array_ops.concat([inputs, h], 1)
        concat = self._linear(args)
        dtype = args.dtype

        i, j, f, o = array_ops.split(value=concat, num_or_size_splits=4, axis=1)
        if self._layer_norm:
            i = self._norm(i, "input", dtype=dtype)
            j = self._norm(j, "transform", dtype=dtype)
            f = self._norm(f, "forget", dtype=dtype)
            o = self._norm(o, "output", dtype=dtype)

        g = self._activation(j)
        if (not isinstance(self._keep_prob, float)) or self._keep_prob < 1:
            g = nn_ops.dropout(g, self._keep_prob, seed=self._seed)

        new_c = (
                c * math_ops.sigmoid(f + self._forget_bias) + math_ops.sigmoid(i) * g)
        if self._layer_norm:
            new_c = self._norm(new_c, "state", dtype=dtype)
        new_h = self._activation(new_c) * math_ops.sigmoid(o)

        new_state = rnn_cell_impl.LSTMStateTuple(new_c, new_h)
        return new_h, new_state


class NASCell(rnn_cell_impl.LayerRNNCell):
    """Neural Architecture Search (NAS) recurrent network cell.

    This implements the recurrent cell from the paper:

      https://arxiv.org/abs/1611.01578

    Barret Zoph and Quoc V. Le.
    "Neural Architecture Search with Reinforcement Learning" Proc. ICLR 2017.

    The class uses an optional projection layer.
    """

    # NAS cell's architecture base.
    _NAS_BASE = 8

    def __init__(self, num_units, num_proj=None, use_bias=False, reuse=None,
                 **kwargs):
        """Initialize the parameters for a NAS cell.

        Args:
          num_units: int, The number of units in the NAS cell.
          num_proj: (optional) int, The output dimensionality for the projection
            matrices.  If None, no projection is performed.
          use_bias: (optional) bool, If True then use biases within the cell. This
            is False by default.
          reuse: (optional) Python boolean describing whether to reuse variables
            in an existing scope.  If not `True`, and the existing scope already has
            the given variables, an error is raised.
          **kwargs: Additional keyword arguments.
        """
        super(NASCell, self).__init__(_reuse=reuse, **kwargs)
        self._num_units = num_units
        self._num_proj = num_proj
        self._use_bias = use_bias
        self._reuse = reuse

        if num_proj is not None:
            self._state_size = rnn_cell_impl.LSTMStateTuple(num_units, num_proj)
            self._output_size = num_proj
        else:
            self._state_size = rnn_cell_impl.LSTMStateTuple(num_units, num_units)
            self._output_size = num_units

    @property
    def state_size(self):
        return self._state_size

    @property
    def output_size(self):
        return self._output_size

    def build(self, inputs_shape):
        input_size = tensor_shape.dimension_value(
            tensor_shape.TensorShape(inputs_shape).with_rank(2)[1])
        if input_size is None:
            raise ValueError("Could not infer input size from inputs.get_shape()[-1]")

        num_proj = self._num_units if self._num_proj is None else self._num_proj

        # Variables for the NAS cell. `recurrent_kernel` is all matrices multiplying
        # the hiddenstate and `kernel` is all matrices multiplying the inputs.
        self.recurrent_kernel = self.add_variable(
            "recurrent_kernel", [num_proj, self._NAS_BASE * self._num_units])
        self.kernel = self.add_variable(
            "kernel", [input_size, self._NAS_BASE * self._num_units])

        if self._use_bias:
            self.bias = self.add_variable("bias",
                                          shape=[self._NAS_BASE * self._num_units],
                                          initializer=init_ops.zeros_initializer)

        # Projection layer if specified
        if self._num_proj is not None:
            self.projection_weights = self.add_variable(
                "projection_weights", [self._num_units, self._num_proj])

        self.built = True

    def call(self, inputs, state):
        """Run one step of NAS Cell.

        Args:
          inputs: input Tensor, 2D, batch x num_units.
          state: This must be a tuple of state Tensors, both `2-D`, with column
            sizes `c_state` and `m_state`.

        Returns:
          A tuple containing:
          - A `2-D, [batch x output_dim]`, Tensor representing the output of the
            NAS Cell after reading `inputs` when previous state was `state`.
            Here output_dim is:
               num_proj if num_proj was set,
               num_units otherwise.
          - Tensor(s) representing the new state of NAS Cell after reading `inputs`
            when the previous state was `state`.  Same type and shape(s) as `state`.

        Raises:
          ValueError: If input size cannot be inferred from inputs via
            static shape inference.
        """
        sigmoid = math_ops.sigmoid
        tanh = math_ops.tanh
        relu = nn_ops.relu

        (c_prev, m_prev) = state

        m_matrix = math_ops.matmul(m_prev, self.recurrent_kernel)
        inputs_matrix = math_ops.matmul(inputs, self.kernel)

        if self._use_bias:
            m_matrix = nn_ops.bias_add(m_matrix, self.bias)

        # The NAS cell branches into 8 different splits for both the hiddenstate
        # and the input
        m_matrix_splits = array_ops.split(
            axis=1, num_or_size_splits=self._NAS_BASE, value=m_matrix)
        inputs_matrix_splits = array_ops.split(
            axis=1, num_or_size_splits=self._NAS_BASE, value=inputs_matrix)

        # First layer
        layer1_0 = sigmoid(inputs_matrix_splits[0] + m_matrix_splits[0])
        layer1_1 = relu(inputs_matrix_splits[1] + m_matrix_splits[1])
        layer1_2 = sigmoid(inputs_matrix_splits[2] + m_matrix_splits[2])
        layer1_3 = relu(inputs_matrix_splits[3] * m_matrix_splits[3])
        layer1_4 = tanh(inputs_matrix_splits[4] + m_matrix_splits[4])
        layer1_5 = sigmoid(inputs_matrix_splits[5] + m_matrix_splits[5])
        layer1_6 = tanh(inputs_matrix_splits[6] + m_matrix_splits[6])
        layer1_7 = sigmoid(inputs_matrix_splits[7] + m_matrix_splits[7])

        # Second layer
        l2_0 = tanh(layer1_0 * layer1_1)
        l2_1 = tanh(layer1_2 + layer1_3)
        l2_2 = tanh(layer1_4 * layer1_5)
        l2_3 = sigmoid(layer1_6 + layer1_7)

        # Inject the cell
        l2_0 = tanh(l2_0 + c_prev)

        # Third layer
        l3_0_pre = l2_0 * l2_1
        new_c = l3_0_pre  # create new cell
        l3_0 = l3_0_pre
        l3_1 = tanh(l2_2 + l2_3)

        # Final layer
        new_m = tanh(l3_0 * l3_1)

        # Projection layer if specified
        if self._num_proj is not None:
            new_m = math_ops.matmul(new_m, self.projection_weights)

        new_state = rnn_cell_impl.LSTMStateTuple(new_c, new_m)
        return new_m, new_state


class UGRNNCell(rnn_cell_impl.RNNCell):
    """Update Gate Recurrent Neural Network (UGRNN) cell.

    Compromise between a LSTM/GRU and a vanilla RNN.  There is only one
    gate, and that is to determine whether the unit should be
    integrating or computing instantaneously.  This is the recurrent
    idea of the feedforward Highway Network.

    This implements the recurrent cell from the paper:

      https://arxiv.org/abs/1611.09913

    Jasmine Collins, Jascha Sohl-Dickstein, and David Sussillo.
    "Capacity and Trainability in Recurrent Neural Networks" Proc. ICLR 2017.
    """

    def __init__(self,
                 num_units,
                 initializer=None,
                 forget_bias=1.0,
                 activation=math_ops.tanh,
                 reuse=None):
        """Initialize the parameters for an UGRNN cell.

        Args:
          num_units: int, The number of units in the UGRNN cell
          initializer: (optional) The initializer to use for the weight matrices.
          forget_bias: (optional) float, default 1.0, The initial bias of the
            forget gate, used to reduce the scale of forgetting at the beginning
            of the training.
          activation: (optional) Activation function of the inner states.
            Default is `tf.tanh`.
          reuse: (optional) Python boolean describing whether to reuse variables
            in an existing scope.  If not `True`, and the existing scope already has
            the given variables, an error is raised.
        """
        super(UGRNNCell, self).__init__(_reuse=reuse)
        self._num_units = num_units
        self._initializer = initializer
        self._forget_bias = forget_bias
        self._activation = activation
        self._reuse = reuse
        self._linear = None

    @property
    def state_size(self):
        return self._num_units

    @property
    def output_size(self):
        return self._num_units

    def call(self, inputs, state):
        """Run one step of UGRNN.

        Args:
          inputs: input Tensor, 2D, batch x input size.
          state: state Tensor, 2D, batch x num units.

        Returns:
          new_output: batch x num units, Tensor representing the output of the UGRNN
            after reading `inputs` when previous state was `state`. Identical to
            `new_state`.
          new_state: batch x num units, Tensor representing the state of the UGRNN
            after reading `inputs` when previous state was `state`.

        Raises:
          ValueError: If input size cannot be inferred from inputs via
            static shape inference.
        """
        sigmoid = math_ops.sigmoid

        input_size = inputs.get_shape().with_rank(2).dims[1]
        if input_size.value is None:
            raise ValueError("Could not infer input size from inputs.get_shape()[-1]")

        with vs.variable_scope(
                vs.get_variable_scope(), initializer=self._initializer):
            cell_inputs = array_ops.concat([inputs, state], 1)
            if self._linear is None:
                self._linear = _Linear(cell_inputs, 2 * self._num_units, True)
            rnn_matrix = self._linear(cell_inputs)

            [g_act, c_act] = array_ops.split(
                axis=1, num_or_size_splits=2, value=rnn_matrix)

            c = self._activation(c_act)
            g = sigmoid(g_act + self._forget_bias)
            new_state = g * state + (1.0 - g) * c
            new_output = new_state

        return new_output, new_state


class IntersectionRNNCell(rnn_cell_impl.RNNCell):
    """Intersection Recurrent Neural Network (+RNN) cell.

    Architecture with coupled recurrent gate as well as coupled depth
    gate, designed to improve information flow through stacked RNNs. As the
    architecture uses depth gating, the dimensionality of the depth
    output (y) also should not change through depth (input size == output size).
    To achieve this, the first layer of a stacked Intersection RNN projects
    the inputs to N (num units) dimensions. Therefore when initializing an
    IntersectionRNNCell, one should set `num_in_proj = N` for the first layer
    and use default settings for subsequent layers.

    This implements the recurrent cell from the paper:

      https://arxiv.org/abs/1611.09913

    Jasmine Collins, Jascha Sohl-Dickstein, and David Sussillo.
    "Capacity and Trainability in Recurrent Neural Networks" Proc. ICLR 2017.

    The Intersection RNN is built for use in deeply stacked
    RNNs so it may not achieve best performance with depth 1.
    """

    def __init__(self,
                 num_units,
                 num_in_proj=None,
                 initializer=None,
                 forget_bias=1.0,
                 y_activation=nn_ops.relu,
                 reuse=None):
        """Initialize the parameters for an +RNN cell.

        Args:
          num_units: int, The number of units in the +RNN cell
          num_in_proj: (optional) int, The input dimensionality for the RNN.
            If creating the first layer of an +RNN, this should be set to
            `num_units`. Otherwise, this should be set to `None` (default).
            If `None`, dimensionality of `inputs` should be equal to `num_units`,
            otherwise ValueError is thrown.
          initializer: (optional) The initializer to use for the weight matrices.
          forget_bias: (optional) float, default 1.0, The initial bias of the
            forget gates, used to reduce the scale of forgetting at the beginning
            of the training.
          y_activation: (optional) Activation function of the states passed
            through depth. Default is 'tf.nn.relu`.
          reuse: (optional) Python boolean describing whether to reuse variables
            in an existing scope.  If not `True`, and the existing scope already has
            the given variables, an error is raised.
        """
        super(IntersectionRNNCell, self).__init__(_reuse=reuse)
        self._num_units = num_units
        self._initializer = initializer
        self._forget_bias = forget_bias
        self._num_input_proj = num_in_proj
        self._y_activation = y_activation
        self._reuse = reuse
        self._linear1 = None
        self._linear2 = None

    @property
    def state_size(self):
        return self._num_units

    @property
    def output_size(self):
        return self._num_units

    def call(self, inputs, state):
        """Run one step of the Intersection RNN.

        Args:
          inputs: input Tensor, 2D, batch x input size.
          state: state Tensor, 2D, batch x num units.

        Returns:
          new_y: batch x num units, Tensor representing the output of the +RNN
            after reading `inputs` when previous state was `state`.
          new_state: batch x num units, Tensor representing the state of the +RNN
            after reading `inputs` when previous state was `state`.

        Raises:
          ValueError: If input size cannot be inferred from `inputs` via
            static shape inference.
          ValueError: If input size != output size (these must be equal when
            using the Intersection RNN).
        """
        sigmoid = math_ops.sigmoid
        tanh = math_ops.tanh

        input_size = inputs.get_shape().with_rank(2).dims[1]
        if input_size.value is None:
            raise ValueError("Could not infer input size from inputs.get_shape()[-1]")

        with vs.variable_scope(
                vs.get_variable_scope(), initializer=self._initializer):
            # read-in projections (should be used for first layer in deep +RNN
            # to transform size of inputs from I --> N)
            if input_size.value != self._num_units:
                if self._num_input_proj:
                    with vs.variable_scope("in_projection"):
                        if self._linear1 is None:
                            self._linear1 = _Linear(inputs, self._num_units, True)
                        inputs = self._linear1(inputs)
                else:
                    raise ValueError("Must have input size == output size for "
                                     "Intersection RNN. To fix, num_in_proj should "
                                     "be set to num_units at cell init.")

            n_dim = i_dim = self._num_units
            cell_inputs = array_ops.concat([inputs, state], 1)
            if self._linear2 is None:
                self._linear2 = _Linear(cell_inputs, 2 * n_dim + 2 * i_dim, True)
            rnn_matrix = self._linear2(cell_inputs)

            gh_act = rnn_matrix[:, :n_dim]  # b x n
            h_act = rnn_matrix[:, n_dim:2 * n_dim]  # b x n
            gy_act = rnn_matrix[:, 2 * n_dim:2 * n_dim + i_dim]  # b x i
            y_act = rnn_matrix[:, 2 * n_dim + i_dim:2 * n_dim + 2 * i_dim]  # b x i

            h = tanh(h_act)
            y = self._y_activation(y_act)
            gh = sigmoid(gh_act + self._forget_bias)
            gy = sigmoid(gy_act + self._forget_bias)

            new_state = gh * state + (1.0 - gh) * h  # passed thru time
            new_y = gy * inputs + (1.0 - gy) * y  # passed thru depth

        return new_y, new_state


_REGISTERED_OPS = None


class CompiledWrapper(rnn_cell_impl.RNNCell):
    """Wraps step execution in an XLA JIT scope."""

    def __init__(self, cell, compile_stateful=False):
        """Create CompiledWrapper cell.

        Args:
          cell: Instance of `RNNCell`.
          compile_stateful: Whether to compile stateful ops like initializers
            and random number generators (default: False).
        """
        self._cell = cell
        self._compile_stateful = compile_stateful

    @property
    def state_size(self):
        return self._cell.state_size

    @property
    def output_size(self):
        return self._cell.output_size

    def zero_state(self, batch_size, dtype):
        with ops.name_scope(type(self).__name__ + "ZeroState", values=[batch_size]):
            return self._cell.zero_state(batch_size, dtype)

    def __call__(self, inputs, state, scope=None):
        if self._compile_stateful:
            compile_ops = True
        else:

            def compile_ops(node_def):
                global _REGISTERED_OPS
                if _REGISTERED_OPS is None:
                    _REGISTERED_OPS = op_def_registry.get_registered_ops()
                return not _REGISTERED_OPS[node_def.op].is_stateful

        with jit.experimental_jit_scope(compile_ops=compile_ops):
            return self._cell(inputs, state, scope=scope)


def _random_exp_initializer(minval, maxval, seed=None, dtype=dtypes.float32):
    """Returns an exponential distribution initializer.

    Args:
      minval: float or a scalar float Tensor. With value > 0. Lower bound of the
          range of random values to generate.
      maxval: float or a scalar float Tensor. With value > minval. Upper bound of
          the range of random values to generate.
      seed: An integer. Used to create random seeds.
      dtype: The data type.

    Returns:
      An initializer that generates tensors with an exponential distribution.
    """

    def _initializer(shape, dtype=dtype, partition_info=None):
        del partition_info  # Unused.
        return math_ops.exp(
            random_ops.random_uniform(
                shape, math_ops.log(minval), math_ops.log(maxval), dtype,
                seed=seed))

    return _initializer


class PhasedLSTMCell(rnn_cell_impl.RNNCell):
    """Phased LSTM recurrent network cell.

    https://arxiv.org/pdf/1610.09513v1.pdf
    """

    def __init__(self,
                 num_units,
                 use_peepholes=False,
                 leak=0.001,
                 ratio_on=0.1,
                 trainable_ratio_on=True,
                 period_init_min=1.0,
                 period_init_max=1000.0,
                 reuse=None):
        """Initialize the Phased LSTM cell.

        Args:
          num_units: int, The number of units in the Phased LSTM cell.
          use_peepholes: bool, set True to enable peephole connections.
          leak: float or scalar float Tensor with value in [0, 1]. Leak applied
              during training.
          ratio_on: float or scalar float Tensor with value in [0, 1]. Ratio of the
              period during which the gates are open.
          trainable_ratio_on: bool, weather ratio_on is trainable.
          period_init_min: float or scalar float Tensor. With value > 0.
              Minimum value of the initialized period.
              The period values are initialized by drawing from the distribution:
              e^U(log(period_init_min), log(period_init_max))
              Where U(.,.) is the uniform distribution.
          period_init_max: float or scalar float Tensor.
              With value > period_init_min. Maximum value of the initialized period.
          reuse: (optional) Python boolean describing whether to reuse variables
            in an existing scope. If not `True`, and the existing scope already has
            the given variables, an error is raised.
        """
        # We pass autocast=False because this layer can accept inputs of different
        # dtypes, so we do not want to automatically cast them to the same dtype.
        super(PhasedLSTMCell, self).__init__(_reuse=reuse, autocast=False)
        self._num_units = num_units
        self._use_peepholes = use_peepholes
        self._leak = leak
        self._ratio_on = ratio_on
        self._trainable_ratio_on = trainable_ratio_on
        self._period_init_min = period_init_min
        self._period_init_max = period_init_max
        self._reuse = reuse
        self._linear1 = None
        self._linear2 = None
        self._linear3 = None

    @property
    def state_size(self):
        return rnn_cell_impl.LSTMStateTuple(self._num_units, self._num_units)

    @property
    def output_size(self):
        return self._num_units

    def _mod(self, x, y):
        """Modulo function that propagates x gradients."""
        return array_ops.stop_gradient(math_ops.mod(x, y) - x) + x

    def _get_cycle_ratio(self, time, phase, period):
        """Compute the cycle ratio in the dtype of the time."""
        phase_casted = math_ops.cast(phase, dtype=time.dtype)
        period_casted = math_ops.cast(period, dtype=time.dtype)
        shifted_time = time - phase_casted
        cycle_ratio = self._mod(shifted_time, period_casted) / period_casted
        return math_ops.cast(cycle_ratio, dtype=dtypes.float32)

    def call(self, inputs, state):
        """Phased LSTM Cell.

        Args:
          inputs: A tuple of 2 Tensor.
             The first Tensor has shape [batch, 1], and type float32 or float64.
             It stores the time.
             The second Tensor has shape [batch, features_size], and type float32.
             It stores the features.
          state: rnn_cell_impl.LSTMStateTuple, state from previous timestep.

        Returns:
          A tuple containing:
          - A Tensor of float32, and shape [batch_size, num_units], representing the
            output of the cell.
          - A rnn_cell_impl.LSTMStateTuple, containing 2 Tensors of float32, shape
            [batch_size, num_units], representing the new state and the output.
        """
        (c_prev, h_prev) = state
        (time, x) = inputs

        in_mask_gates = [x, h_prev]
        if self._use_peepholes:
            in_mask_gates.append(c_prev)

        with vs.variable_scope("mask_gates"):
            if self._linear1 is None:
                self._linear1 = _Linear(in_mask_gates, 2 * self._num_units, True)

            mask_gates = math_ops.sigmoid(self._linear1(in_mask_gates))
            [input_gate, forget_gate] = array_ops.split(
                axis=1, num_or_size_splits=2, value=mask_gates)

        with vs.variable_scope("new_input"):
            if self._linear2 is None:
                self._linear2 = _Linear([x, h_prev], self._num_units, True)
            new_input = math_ops.tanh(self._linear2([x, h_prev]))

        new_c = (c_prev * forget_gate + input_gate * new_input)

        in_out_gate = [x, h_prev]
        if self._use_peepholes:
            in_out_gate.append(new_c)

        with vs.variable_scope("output_gate"):
            if self._linear3 is None:
                self._linear3 = _Linear(in_out_gate, self._num_units, True)
            output_gate = math_ops.sigmoid(self._linear3(in_out_gate))

        new_h = math_ops.tanh(new_c) * output_gate

        period = vs.get_variable(
            "period", [self._num_units],
            initializer=_random_exp_initializer(self._period_init_min,
                                                self._period_init_max))
        phase = vs.get_variable(
            "phase", [self._num_units],
            initializer=init_ops.random_uniform_initializer(0.,
                                                            period.initial_value))
        ratio_on = vs.get_variable(
            "ratio_on", [self._num_units],
            initializer=init_ops.constant_initializer(self._ratio_on),
            trainable=self._trainable_ratio_on)

        cycle_ratio = self._get_cycle_ratio(time, phase, period)

        k_up = 2 * cycle_ratio / ratio_on
        k_down = 2 - k_up
        k_closed = self._leak * cycle_ratio

        k = array_ops.where(cycle_ratio < ratio_on, k_down, k_closed)
        k = array_ops.where(cycle_ratio < 0.5 * ratio_on, k_up, k)

        new_c = k * new_c + (1 - k) * c_prev
        new_h = k * new_h + (1 - k) * h_prev

        new_state = rnn_cell_impl.LSTMStateTuple(new_c, new_h)

        return new_h, new_state


class ConvLSTMCell(rnn_cell_impl.RNNCell):
    """Convolutional LSTM recurrent network cell.

    https://arxiv.org/pdf/1506.04214v1.pdf
    """

    def __init__(self,
                 conv_ndims,
                 input_shape,
                 output_channels,
                 kernel_shape,
                 use_bias=True,
                 skip_connection=False,
                 forget_bias=1.0,
                 initializers=None,
                 name="conv_lstm_cell"):
        """Construct ConvLSTMCell.

        Args:
          conv_ndims: Convolution dimensionality (1, 2 or 3).
          input_shape: Shape of the input as int tuple, excluding the batch size.
          output_channels: int, number of output channels of the conv LSTM.
          kernel_shape: Shape of kernel as an int tuple (of size 1, 2 or 3).
          use_bias: (bool) Use bias in convolutions.
          skip_connection: If set to `True`, concatenate the input to the
            output of the conv LSTM. Default: `False`.
          forget_bias: Forget bias.
          initializers: Unused.
          name: Name of the module.

        Raises:
          ValueError: If `skip_connection` is `True` and stride is different from 1
            or if `input_shape` is incompatible with `conv_ndims`.
        """
        super(ConvLSTMCell, self).__init__(name=name)

        if conv_ndims != len(input_shape) - 1:
            raise ValueError("Invalid input_shape {} for conv_ndims={}.".format(
                input_shape, conv_ndims))

        self._conv_ndims = conv_ndims
        self._input_shape = input_shape
        self._output_channels = output_channels
        self._kernel_shape = list(kernel_shape)
        self._use_bias = use_bias
        self._forget_bias = forget_bias
        self._skip_connection = skip_connection

        self._total_output_channels = output_channels
        if self._skip_connection:
            self._total_output_channels += self._input_shape[-1]

        state_size = tensor_shape.TensorShape(
            self._input_shape[:-1] + [self._output_channels])
        self._state_size = rnn_cell_impl.LSTMStateTuple(state_size, state_size)
        self._output_size = tensor_shape.TensorShape(
            self._input_shape[:-1] + [self._total_output_channels])

    @property
    def output_size(self):
        return self._output_size

    @property
    def state_size(self):
        return self._state_size

    def call(self, inputs, state, scope=None):
        cell, hidden = state
        new_hidden = _conv([inputs, hidden], self._kernel_shape,
                           4 * self._output_channels, self._use_bias)
        gates = array_ops.split(
            value=new_hidden, num_or_size_splits=4, axis=self._conv_ndims + 1)

        input_gate, new_input, forget_gate, output_gate = gates
        new_cell = math_ops.sigmoid(forget_gate + self._forget_bias) * cell
        new_cell += math_ops.sigmoid(input_gate) * math_ops.tanh(new_input)
        output = math_ops.tanh(new_cell) * math_ops.sigmoid(output_gate)

        if self._skip_connection:
            output = array_ops.concat([output, inputs], axis=-1)
        new_state = rnn_cell_impl.LSTMStateTuple(new_cell, output)
        return output, new_state


class Conv1DLSTMCell(ConvLSTMCell):
    """1D Convolutional LSTM recurrent network cell.

    https://arxiv.org/pdf/1506.04214v1.pdf
    """

    def __init__(self, name="conv_1d_lstm_cell", **kwargs):
        """Construct Conv1DLSTM. See `ConvLSTMCell` for more details."""
        super(Conv1DLSTMCell, self).__init__(conv_ndims=1, name=name, **kwargs)


class Conv2DLSTMCell(ConvLSTMCell):
    """2D Convolutional LSTM recurrent network cell.

    https://arxiv.org/pdf/1506.04214v1.pdf
    """

    def __init__(self, name="conv_2d_lstm_cell", **kwargs):
        """Construct Conv2DLSTM. See `ConvLSTMCell` for more details."""
        super(Conv2DLSTMCell, self).__init__(conv_ndims=2, name=name, **kwargs)


class Conv3DLSTMCell(ConvLSTMCell):
    """3D Convolutional LSTM recurrent network cell.

    https://arxiv.org/pdf/1506.04214v1.pdf
    """

    def __init__(self, name="conv_3d_lstm_cell", **kwargs):
        """Construct Conv3DLSTM. See `ConvLSTMCell` for more details."""
        super(Conv3DLSTMCell, self).__init__(conv_ndims=3, name=name, **kwargs)


def _conv(args, filter_size, num_features, bias, bias_start=0.0):
    """Convolution.

    Args:
      args: a Tensor or a list of Tensors of dimension 3D, 4D or 5D,
      batch x n, Tensors.
      filter_size: int tuple of filter shape (of size 1, 2 or 3).
      num_features: int, number of features.
      bias: Whether to use biases in the convolution layer.
      bias_start: starting value to initialize the bias; 0 by default.

    Returns:
      A 3D, 4D, or 5D Tensor with shape [batch ... num_features]

    Raises:
      ValueError: if some of the arguments has unspecified or wrong shape.
    """

    # Calculate the total size of arguments on dimension 1.
    total_arg_size_depth = 0
    shapes = [a.get_shape().as_list() for a in args]
    shape_length = len(shapes[0])
    for shape in shapes:
        if len(shape) not in [3, 4, 5]:
            raise ValueError("Conv Linear expects 3D, 4D "
                             "or 5D arguments: %s" % str(shapes))
        if len(shape) != len(shapes[0]):
            raise ValueError("Conv Linear expects all args "
                             "to be of same Dimension: %s" % str(shapes))
        else:
            total_arg_size_depth += shape[-1]
    dtype = [a.dtype for a in args][0]

    # determine correct conv operation
    if shape_length == 3:
        conv_op = nn_ops.conv1d
        strides = 1
    elif shape_length == 4:
        conv_op = nn_ops.conv2d
        strides = shape_length * [1]
    elif shape_length == 5:
        conv_op = nn_ops.conv3d
        strides = shape_length * [1]

    # Now the computation.
    kernel = vs.get_variable(
        "kernel", filter_size + [total_arg_size_depth, num_features], dtype=dtype)
    if len(args) == 1:
        res = conv_op(args[0], kernel, strides, padding="SAME")
    else:
        res = conv_op(
            array_ops.concat(axis=shape_length - 1, values=args),
            kernel,
            strides,
            padding="SAME")
    if not bias:
        return res
    bias_term = vs.get_variable(
        "biases", [num_features],
        dtype=dtype,
        initializer=init_ops.constant_initializer(bias_start, dtype=dtype))
    return res + bias_term


class GLSTMCell(rnn_cell_impl.RNNCell):
    """Group LSTM cell (G-LSTM).

    The implementation is based on:

      https://arxiv.org/abs/1703.10722

    O. Kuchaiev and B. Ginsburg
    "Factorization Tricks for LSTM Networks", ICLR 2017 workshop.

    In brief, a G-LSTM cell consists of one LSTM sub-cell per group, where each
    sub-cell operates on an evenly-sized sub-vector of the input and produces an
    evenly-sized sub-vector of the output.  For example, a G-LSTM cell with 128
    units and 4 groups consists of 4 LSTMs sub-cells with 32 units each.  If that
    G-LSTM cell is fed a 200-dim input, then each sub-cell receives a 50-dim part
    of the input and produces a 32-dim part of the output.
    """

    def __init__(self,
                 num_units,
                 initializer=None,
                 num_proj=None,
                 number_of_groups=1,
                 forget_bias=1.0,
                 activation=math_ops.tanh,
                 reuse=None):
        """Initialize the parameters of G-LSTM cell.

        Args:
          num_units: int, The number of units in the G-LSTM cell
          initializer: (optional) The initializer to use for the weight and
            projection matrices.
          num_proj: (optional) int, The output dimensionality for the projection
            matrices.  If None, no projection is performed.
          number_of_groups: (optional) int, number of groups to use.
            If `number_of_groups` is 1, then it should be equivalent to LSTM cell
          forget_bias: Biases of the forget gate are initialized by default to 1
            in order to reduce the scale of forgetting at the beginning of
            the training.
          activation: Activation function of the inner states.
          reuse: (optional) Python boolean describing whether to reuse variables
            in an existing scope.  If not `True`, and the existing scope already
            has the given variables, an error is raised.

        Raises:
          ValueError: If `num_units` or `num_proj` is not divisible by
            `number_of_groups`.
        """
        super(GLSTMCell, self).__init__(_reuse=reuse)
        self._num_units = num_units
        self._initializer = initializer
        self._num_proj = num_proj
        self._forget_bias = forget_bias
        self._activation = activation
        self._number_of_groups = number_of_groups

        if self._num_units % self._number_of_groups != 0:
            raise ValueError("num_units must be divisible by number_of_groups")
        if self._num_proj:
            if self._num_proj % self._number_of_groups != 0:
                raise ValueError("num_proj must be divisible by number_of_groups")
            self._group_shape = [
                int(self._num_proj / self._number_of_groups),
                int(self._num_units / self._number_of_groups)
            ]
        else:
            self._group_shape = [
                int(self._num_units / self._number_of_groups),
                int(self._num_units / self._number_of_groups)
            ]

        if num_proj:
            self._state_size = rnn_cell_impl.LSTMStateTuple(num_units, num_proj)
            self._output_size = num_proj
        else:
            self._state_size = rnn_cell_impl.LSTMStateTuple(num_units, num_units)
            self._output_size = num_units
        self._linear1 = [None] * number_of_groups
        self._linear2 = None

    @property
    def state_size(self):
        return self._state_size

    @property
    def output_size(self):
        return self._output_size

    def _get_input_for_group(self, inputs, group_id, group_size):
        """Slices inputs into groups to prepare for processing by cell's groups.

        Args:
          inputs: cell input or it's previous state,
                  a Tensor, 2D, [batch x num_units]
          group_id: group id, a Scalar, for which to prepare input
          group_size: size of the group

        Returns:
          subset of inputs corresponding to group "group_id",
          a Tensor, 2D, [batch x num_units/number_of_groups]
        """
        return array_ops.slice(
            input_=inputs,
            begin=[0, group_id * group_size],
            size=[self._batch_size, group_size],
            name=("GLSTM_group%d_input_generation" % group_id))

    def call(self, inputs, state):
        """Run one step of G-LSTM.

        Args:
          inputs: input Tensor, 2D, [batch x num_inputs].  num_inputs must be
            statically-known and evenly divisible into groups.  The innermost
            vectors of the inputs are split into evenly-sized sub-vectors and fed
            into the per-group LSTM sub-cells.
          state: this must be a tuple of state Tensors, both `2-D`, with column
            sizes `c_state` and `m_state`.

        Returns:
          A tuple containing:

          - A `2-D, [batch x output_dim]`, Tensor representing the output of the
            G-LSTM after reading `inputs` when previous state was `state`.
            Here output_dim is:
               num_proj if num_proj was set,
               num_units otherwise.
          - LSTMStateTuple representing the new state of G-LSTM cell
            after reading `inputs` when the previous state was `state`.

        Raises:
          ValueError: If input size cannot be inferred from inputs via
            static shape inference, or if the input shape is incompatible
            with the number of groups.
        """
        (c_prev, m_prev) = state

        self._batch_size = tensor_shape.dimension_value(
            inputs.shape[0]) or array_ops.shape(inputs)[0]

        # If the input size is statically-known, calculate and validate its group
        # size.  Otherwise, use the output group size.
        input_size = tensor_shape.dimension_value(inputs.shape[1])
        if input_size is None:
            raise ValueError("input size must be statically known")
        if input_size % self._number_of_groups != 0:
            raise ValueError(
                "input size (%d) must be divisible by number_of_groups (%d)" %
                (input_size, self._number_of_groups))
        input_group_size = int(input_size / self._number_of_groups)

        dtype = inputs.dtype
        scope = vs.get_variable_scope()
        with vs.variable_scope(scope, initializer=self._initializer):
            i_parts = []
            j_parts = []
            f_parts = []
            o_parts = []

            for group_id in range(self._number_of_groups):
                with vs.variable_scope("group%d" % group_id):
                    x_g_id = array_ops.concat(
                        [
                            self._get_input_for_group(inputs, group_id, input_group_size),
                            self._get_input_for_group(m_prev, group_id,
                                                      self._group_shape[0])
                        ],
                        axis=1)
                    linear = self._linear1[group_id]
                    if linear is None:
                        linear = _Linear(x_g_id, 4 * self._group_shape[1], False)
                        self._linear1[group_id] = linear
                    R_k = linear(x_g_id)  # pylint: disable=invalid-name
                    i_k, j_k, f_k, o_k = array_ops.split(R_k, 4, 1)

                i_parts.append(i_k)
                j_parts.append(j_k)
                f_parts.append(f_k)
                o_parts.append(o_k)

            bi = vs.get_variable(
                name="bias_i",
                shape=[self._num_units],
                dtype=dtype,
                initializer=init_ops.constant_initializer(0.0, dtype=dtype))
            bj = vs.get_variable(
                name="bias_j",
                shape=[self._num_units],
                dtype=dtype,
                initializer=init_ops.constant_initializer(0.0, dtype=dtype))
            bf = vs.get_variable(
                name="bias_f",
                shape=[self._num_units],
                dtype=dtype,
                initializer=init_ops.constant_initializer(0.0, dtype=dtype))
            bo = vs.get_variable(
                name="bias_o",
                shape=[self._num_units],
                dtype=dtype,
                initializer=init_ops.constant_initializer(0.0, dtype=dtype))

            i = nn_ops.bias_add(array_ops.concat(i_parts, axis=1), bi)
            j = nn_ops.bias_add(array_ops.concat(j_parts, axis=1), bj)
            f = nn_ops.bias_add(array_ops.concat(f_parts, axis=1), bf)
            o = nn_ops.bias_add(array_ops.concat(o_parts, axis=1), bo)

        c = (
                math_ops.sigmoid(f + self._forget_bias) * c_prev +
                math_ops.sigmoid(i) * math_ops.tanh(j))
        m = math_ops.sigmoid(o) * self._activation(c)

        if self._num_proj is not None:
            with vs.variable_scope("projection"):
                if self._linear2 is None:
                    self._linear2 = _Linear(m, self._num_proj, False)
                m = self._linear2(m)

        new_state = rnn_cell_impl.LSTMStateTuple(c, m)
        return m, new_state


class LayerNormLSTMCell(rnn_cell_impl.RNNCell):
    """Long short-term memory unit (LSTM) recurrent network cell.

    The default non-peephole implementation is based on:

      https://pdfs.semanticscholar.org/1154/0131eae85b2e11d53df7f1360eeb6476e7f4.pdf

    Felix Gers, Jurgen Schmidhuber, and Fred Cummins.
    "Learning to forget: Continual prediction with LSTM." IET, 850-855, 1999.

    The peephole implementation is based on:

      https://research.google.com/pubs/archive/43905.pdf

    Hasim Sak, Andrew Senior, and Francoise Beaufays.
    "Long short-term memory recurrent neural network architectures for
     large scale acoustic modeling." INTERSPEECH, 2014.

    The class uses optional peep-hole connections, optional cell clipping, and
    an optional projection layer.

    Layer normalization implementation is based on:

      https://arxiv.org/abs/1607.06450.

    "Layer Normalization"
    Jimmy Lei Ba, Jamie Ryan Kiros, Geoffrey E. Hinton

    and is applied before the internal nonlinearities.

    """

    def __init__(self,
                 num_units,
                 use_peepholes=False,
                 cell_clip=None,
                 initializer=None,
                 num_proj=None,
                 proj_clip=None,
                 forget_bias=1.0,
                 activation=None,
                 layer_norm=False,
                 norm_gain=1.0,
                 norm_shift=0.0,
                 reuse=None):
        """Initialize the parameters for an LSTM cell.

        Args:
          num_units: int, The number of units in the LSTM cell
          use_peepholes: bool, set True to enable diagonal/peephole connections.
          cell_clip: (optional) A float value, if provided the cell state is clipped
            by this value prior to the cell output activation.
          initializer: (optional) The initializer to use for the weight and
            projection matrices.
          num_proj: (optional) int, The output dimensionality for the projection
            matrices.  If None, no projection is performed.
          proj_clip: (optional) A float value.  If `num_proj > 0` and `proj_clip` is
            provided, then the projected values are clipped elementwise to within
            `[-proj_clip, proj_clip]`.
          forget_bias: Biases of the forget gate are initialized by default to 1
            in order to reduce the scale of forgetting at the beginning of
            the training. Must set it manually to `0.0` when restoring from
            CudnnLSTM trained checkpoints.
          activation: Activation function of the inner states.  Default: `tanh`.
          layer_norm: If `True`, layer normalization will be applied.
          norm_gain: float, The layer normalization gain initial value. If
            `layer_norm` has been set to `False`, this argument will be ignored.
          norm_shift: float, The layer normalization shift initial value. If
            `layer_norm` has been set to `False`, this argument will be ignored.
          reuse: (optional) Python boolean describing whether to reuse variables
            in an existing scope.  If not `True`, and the existing scope already has
            the given variables, an error is raised.

          When restoring from CudnnLSTM-trained checkpoints, must use
          CudnnCompatibleLSTMCell instead.
        """
        super(LayerNormLSTMCell, self).__init__(_reuse=reuse)

        self._num_units = num_units
        self._use_peepholes = use_peepholes
        self._cell_clip = cell_clip
        self._initializer = initializer
        self._num_proj = num_proj
        self._proj_clip = proj_clip
        self._forget_bias = forget_bias
        self._activation = activation or math_ops.tanh
        self._layer_norm = layer_norm
        self._norm_gain = norm_gain
        self._norm_shift = norm_shift

        if num_proj:
            self._state_size = (rnn_cell_impl.LSTMStateTuple(num_units, num_proj))
            self._output_size = num_proj
        else:
            self._state_size = (rnn_cell_impl.LSTMStateTuple(num_units, num_units))
            self._output_size = num_units

    @property
    def state_size(self):
        return self._state_size

    @property
    def output_size(self):
        return self._output_size

    def _linear(self,
                args,
                output_size,
                bias,
                bias_initializer=None,
                kernel_initializer=None,
                layer_norm=False):
        """Linear map: sum_i(args[i] * W[i]), where W[i] is a Variable.

        Args:
          args: a 2D Tensor or a list of 2D, batch x n, Tensors.
          output_size: int, second dimension of W[i].
          bias: boolean, whether to add a bias term or not.
          bias_initializer: starting value to initialize the bias
            (default is all zeros).
          kernel_initializer: starting value to initialize the weight.
          layer_norm: boolean, whether to apply layer normalization.


        Returns:
          A 2D Tensor with shape [batch x output_size] taking value
          sum_i(args[i] * W[i]), where each W[i] is a newly created Variable.

        Raises:
          ValueError: if some of the arguments has unspecified or wrong shape.
        """
        if args is None or (nest.is_sequence(args) and not args):
            raise ValueError("`args` must be specified")
        if not nest.is_sequence(args):
            args = [args]

        # Calculate the total size of arguments on dimension 1.
        total_arg_size = 0
        shapes = [a.get_shape() for a in args]
        for shape in shapes:
            if shape.ndims != 2:
                raise ValueError("linear is expecting 2D arguments: %s" % shapes)
            if tensor_shape.dimension_value(shape[1]) is None:
                raise ValueError("linear expects shape[1] to be provided for shape %s, "
                                 "but saw %s" % (shape, shape[1]))
            else:
                total_arg_size += tensor_shape.dimension_value(shape[1])

        dtype = [a.dtype for a in args][0]

        # Now the computation.
        scope = vs.get_variable_scope()
        with vs.variable_scope(scope) as outer_scope:
            weights = vs.get_variable(
                "kernel", [total_arg_size, output_size],
                dtype=dtype,
                initializer=kernel_initializer)
            if len(args) == 1:
                res = math_ops.matmul(args[0], weights)
            else:
                res = math_ops.matmul(array_ops.concat(args, 1), weights)
            if not bias:
                return res
            with vs.variable_scope(outer_scope) as inner_scope:
                inner_scope.set_partitioner(None)
                if bias_initializer is None:
                    bias_initializer = init_ops.constant_initializer(0.0, dtype=dtype)
                biases = vs.get_variable(
                    "bias", [output_size], dtype=dtype, initializer=bias_initializer)

        if not layer_norm:
            res = nn_ops.bias_add(res, biases)

        return res

    def call(self, inputs, state):
        """Run one step of LSTM.

        Args:
          inputs: input Tensor, 2D, batch x num_units.
          state: this must be a tuple of state Tensors,
           both `2-D`, with column sizes `c_state` and
            `m_state`.

        Returns:
          A tuple containing:

          - A `2-D, [batch x output_dim]`, Tensor representing the output of the
            LSTM after reading `inputs` when previous state was `state`.
            Here output_dim is:
               num_proj if num_proj was set,
               num_units otherwise.
          - Tensor(s) representing the new state of LSTM after reading `inputs` when
            the previous state was `state`.  Same type and shape(s) as `state`.

        Raises:
          ValueError: If input size cannot be inferred from inputs via
            static shape inference.
        """
        sigmoid = math_ops.sigmoid

        (c_prev, m_prev) = state

        dtype = inputs.dtype
        input_size = inputs.get_shape().with_rank(2).dims[1]
        if input_size.value is None:
            raise ValueError("Could not infer input size from inputs.get_shape()[-1]")
        scope = vs.get_variable_scope()
        with vs.variable_scope(scope, initializer=self._initializer) as unit_scope:

            # i = input_gate, j = new_input, f = forget_gate, o = output_gate
            lstm_matrix = self._linear(
                [inputs, m_prev],
                4 * self._num_units,
                bias=True,
                bias_initializer=None,
                layer_norm=self._layer_norm)
            i, j, f, o = array_ops.split(
                value=lstm_matrix, num_or_size_splits=4, axis=1)

            if self._layer_norm:
                i = _norm(self._norm_gain, self._norm_shift, i, "input")
                j = _norm(self._norm_gain, self._norm_shift, j, "transform")
                f = _norm(self._norm_gain, self._norm_shift, f, "forget")
                o = _norm(self._norm_gain, self._norm_shift, o, "output")

            # Diagonal connections
            if self._use_peepholes:
                with vs.variable_scope(unit_scope):
                    w_f_diag = vs.get_variable(
                        "w_f_diag", shape=[self._num_units], dtype=dtype)
                    w_i_diag = vs.get_variable(
                        "w_i_diag", shape=[self._num_units], dtype=dtype)
                    w_o_diag = vs.get_variable(
                        "w_o_diag", shape=[self._num_units], dtype=dtype)

            if self._use_peepholes:
                c = (
                        sigmoid(f + self._forget_bias + w_f_diag * c_prev) * c_prev +
                        sigmoid(i + w_i_diag * c_prev) * self._activation(j))
            else:
                c = (
                        sigmoid(f + self._forget_bias) * c_prev +
                        sigmoid(i) * self._activation(j))

            if self._layer_norm:
                c = _norm(self._norm_gain, self._norm_shift, c, "state")

            if self._cell_clip is not None:
                # pylint: disable=invalid-unary-operand-type
                c = clip_ops.clip_by_value(c, -self._cell_clip, self._cell_clip)
                # pylint: enable=invalid-unary-operand-type
            if self._use_peepholes:
                m = sigmoid(o + w_o_diag * c) * self._activation(c)
            else:
                m = sigmoid(o) * self._activation(c)

            if self._num_proj is not None:
                with vs.variable_scope("projection"):
                    m = self._linear(m, self._num_proj, bias=False)

                if self._proj_clip is not None:
                    # pylint: disable=invalid-unary-operand-type
                    m = clip_ops.clip_by_value(m, -self._proj_clip, self._proj_clip)
                    # pylint: enable=invalid-unary-operand-type

        new_state = (rnn_cell_impl.LSTMStateTuple(c, m))
        return m, new_state


class SRUCell(rnn_cell_impl.LayerRNNCell):
    """SRU, Simple Recurrent Unit.

       Implementation based on
       Training RNNs as Fast as CNNs (cf. https://arxiv.org/abs/1709.02755).

       This variation of RNN cell is characterized by the simplified data
       dependence
       between hidden states of two consecutive time steps. Traditionally, hidden
       states from a cell at time step t-1 needs to be multiplied with a matrix
       W_hh before being fed into the ensuing cell at time step t.
       This flavor of RNN replaces the matrix multiplication between h_{t-1}
       and W_hh with a pointwise multiplication, resulting in performance
       gain.

    Args:
      num_units: int, The number of units in the SRU cell.
      activation: Nonlinearity to use.  Default: `tanh`.
      reuse: (optional) Python boolean describing whether to reuse variables
        in an existing scope.  If not `True`, and the existing scope already has
        the given variables, an error is raised.
      name: (optional) String, the name of the layer. Layers with the same name
        will share weights, but to avoid mistakes we require reuse=True in such
        cases.
      **kwargs: Additional keyword arguments.
    """

    def __init__(self, num_units, activation=None, reuse=None, name=None,
                 **kwargs):
        super(SRUCell, self).__init__(_reuse=reuse, name=name, **kwargs)
        self._num_units = num_units
        self._activation = activation or math_ops.tanh

        # Restrict inputs to be 2-dimensional matrices
        self.input_spec = input_spec.InputSpec(ndim=2)

    @property
    def state_size(self):
        return self._num_units

    @property
    def output_size(self):
        return self._num_units

    def build(self, inputs_shape):
        if tensor_shape.dimension_value(inputs_shape[1]) is None:
            raise ValueError(
                "Expected inputs.shape[-1] to be known, saw shape: %s" % inputs_shape)

        input_depth = tensor_shape.dimension_value(inputs_shape[1])

        # pylint: disable=protected-access
        self._kernel = self.add_variable(
            rnn_cell_impl._WEIGHTS_VARIABLE_NAME,
            shape=[input_depth, 4 * self._num_units])
        # pylint: enable=protected-access
        self._bias = self.add_variable(
            rnn_cell_impl._BIAS_VARIABLE_NAME,  # pylint: disable=protected-access
            shape=[2 * self._num_units],
            initializer=init_ops.zeros_initializer)

        self._built = True

    def call(self, inputs, state):
        """Simple recurrent unit (SRU) with num_units cells."""

        U = math_ops.matmul(inputs, self._kernel)  # pylint: disable=invalid-name
        x_bar, f_intermediate, r_intermediate, x_tx = array_ops.split(
            value=U, num_or_size_splits=4, axis=1)

        f_r = math_ops.sigmoid(
            nn_ops.bias_add(
                array_ops.concat([f_intermediate, r_intermediate], 1), self._bias))
        f, r = array_ops.split(value=f_r, num_or_size_splits=2, axis=1)

        c = f * state + (1.0 - f) * x_bar
        h = r * self._activation(c) + (1.0 - r) * x_tx

        return h, c


class WeightNormLSTMCell(rnn_cell_impl.RNNCell):
    """Weight normalized LSTM Cell. Adapted from `rnn_cell_impl.LSTMCell`.

      The weight-norm implementation is based on:
      https://arxiv.org/abs/1602.07868
      Tim Salimans, Diederik P. Kingma.
      Weight Normalization: A Simple Reparameterization to Accelerate
      Training of Deep Neural Networks

      The default LSTM implementation based on:

        https://pdfs.semanticscholar.org/1154/0131eae85b2e11d53df7f1360eeb6476e7f4.pdf

      Felix Gers, Jurgen Schmidhuber, and Fred Cummins.
      "Learning to forget: Continual prediction with LSTM." IET, 850-855, 1999.

      The class uses optional peephole connections, optional cell clipping
      and an optional projection layer.

      The optional peephole implementation is based on:
      https://research.google.com/pubs/archive/43905.pdf
      Hasim Sak, Andrew Senior, and Francoise Beaufays.
      "Long short-term memory recurrent neural network architectures for
      large scale acoustic modeling." INTERSPEECH, 2014.
    """

    def __init__(self,
                 num_units,
                 norm=True,
                 use_peepholes=False,
                 cell_clip=None,
                 initializer=None,
                 num_proj=None,
                 proj_clip=None,
                 forget_bias=1,
                 activation=None,
                 reuse=None):
        """Initialize the parameters of a weight-normalized LSTM cell.

        Args:
          num_units: int, The number of units in the LSTM cell
          norm: If `True`, apply normalization to the weight matrices. If False,
            the result is identical to that obtained from `rnn_cell_impl.LSTMCell`
          use_peepholes: bool, set `True` to enable diagonal/peephole connections.
          cell_clip: (optional) A float value, if provided the cell state is clipped
            by this value prior to the cell output activation.
          initializer: (optional) The initializer to use for the weight matrices.
          num_proj: (optional) int, The output dimensionality for the projection
            matrices.  If None, no projection is performed.
          proj_clip: (optional) A float value.  If `num_proj > 0` and `proj_clip` is
            provided, then the projected values are clipped elementwise to within
            `[-proj_clip, proj_clip]`.
          forget_bias: Biases of the forget gate are initialized by default to 1
            in order to reduce the scale of forgetting at the beginning of
            the training.
          activation: Activation function of the inner states.  Default: `tanh`.
          reuse: (optional) Python boolean describing whether to reuse variables
            in an existing scope.  If not `True`, and the existing scope already has
            the given variables, an error is raised.
        """
        super(WeightNormLSTMCell, self).__init__(_reuse=reuse)

        self._scope = "wn_lstm_cell"
        self._num_units = num_units
        self._norm = norm
        self._initializer = initializer
        self._use_peepholes = use_peepholes
        self._cell_clip = cell_clip
        self._num_proj = num_proj
        self._proj_clip = proj_clip
        self._activation = activation or math_ops.tanh
        self._forget_bias = forget_bias

        self._weights_variable_name = "kernel"
        self._bias_variable_name = "bias"

        if num_proj:
            self._state_size = rnn_cell_impl.LSTMStateTuple(num_units, num_proj)
            self._output_size = num_proj
        else:
            self._state_size = rnn_cell_impl.LSTMStateTuple(num_units, num_units)
            self._output_size = num_units

    @property
    def state_size(self):
        return self._state_size

    @property
    def output_size(self):
        return self._output_size

    def _normalize(self, weight, name):
        """Apply weight normalization.

        Args:
          weight: a 2D tensor with known number of columns.
          name: string, variable name for the normalizer.
        Returns:
          A tensor with the same shape as `weight`.
        """

        output_size = weight.get_shape().as_list()[1]
        g = vs.get_variable(name, [output_size], dtype=weight.dtype)
        return nn_impl.l2_normalize(weight, axis=0) * g

    def _linear(self,
                args,
                output_size,
                norm,
                bias,
                bias_initializer=None,
                kernel_initializer=None):
        """Linear map: sum_i(args[i] * W[i]), where W[i] is a variable.

        Args:
          args: a 2D Tensor or a list of 2D, batch x n, Tensors.
          output_size: int, second dimension of W[i].
          norm: bool, whether to normalize the weights.
          bias: boolean, whether to add a bias term or not.
          bias_initializer: starting value to initialize the bias
            (default is all zeros).
          kernel_initializer: starting value to initialize the weight.

        Returns:
          A 2D Tensor with shape [batch x output_size] equal to
          sum_i(args[i] * W[i]), where W[i]s are newly created matrices.

        Raises:
          ValueError: if some of the arguments has unspecified or wrong shape.
        """
        if args is None or (nest.is_sequence(args) and not args):
            raise ValueError("`args` must be specified")
        if not nest.is_sequence(args):
            args = [args]

        # Calculate the total size of arguments on dimension 1.
        total_arg_size = 0
        shapes = [a.get_shape() for a in args]
        for shape in shapes:
            if shape.ndims != 2:
                raise ValueError("linear is expecting 2D arguments: %s" % shapes)
            if tensor_shape.dimension_value(shape[1]) is None:
                raise ValueError("linear expects shape[1] to be provided for shape %s, "
                                 "but saw %s" % (shape, shape[1]))
            else:
                total_arg_size += tensor_shape.dimension_value(shape[1])

        dtype = [a.dtype for a in args][0]

        # Now the computation.
        scope = vs.get_variable_scope()
        with vs.variable_scope(scope) as outer_scope:
            weights = vs.get_variable(
                self._weights_variable_name, [total_arg_size, output_size],
                dtype=dtype,
                initializer=kernel_initializer)
            if norm:
                wn = []
                st = 0
                with ops.control_dependencies(None):
                    for i in range(len(args)):
                        en = st + tensor_shape.dimension_value(shapes[i][1])
                        wn.append(
                            self._normalize(weights[st:en, :], name="norm_{}".format(i)))
                        st = en

                    weights = array_ops.concat(wn, axis=0)

            if len(args) == 1:
                res = math_ops.matmul(args[0], weights)
            else:
                res = math_ops.matmul(array_ops.concat(args, 1), weights)
            if not bias:
                return res

            with vs.variable_scope(outer_scope) as inner_scope:
                inner_scope.set_partitioner(None)
                if bias_initializer is None:
                    bias_initializer = init_ops.constant_initializer(0.0, dtype=dtype)

                biases = vs.get_variable(
                    self._bias_variable_name, [output_size],
                    dtype=dtype,
                    initializer=bias_initializer)

            return nn_ops.bias_add(res, biases)

    def call(self, inputs, state):
        """Run one step of LSTM.

        Args:
          inputs: input Tensor, 2D, batch x num_units.
          state: A tuple of state Tensors, both `2-D`, with column sizes
           `c_state` and `m_state`.

        Returns:
          A tuple containing:

          - A `2-D, [batch x output_dim]`, Tensor representing the output of the
            LSTM after reading `inputs` when previous state was `state`.
            Here output_dim is:
               num_proj if num_proj was set,
               num_units otherwise.
          - Tensor(s) representing the new state of LSTM after reading `inputs` when
            the previous state was `state`.  Same type and shape(s) as `state`.

        Raises:
          ValueError: If input size cannot be inferred from inputs via
            static shape inference.
        """
        dtype = inputs.dtype
        num_units = self._num_units
        sigmoid = math_ops.sigmoid
        c, h = state

        input_size = inputs.get_shape().with_rank(2).dims[1]
        if input_size.value is None:
            raise ValueError("Could not infer input size from inputs.get_shape()[-1]")

        with vs.variable_scope(self._scope, initializer=self._initializer):

            concat = self._linear(
                [inputs, h], 4 * num_units, norm=self._norm, bias=True)

            # i = input_gate, j = new_input, f = forget_gate, o = output_gate
            i, j, f, o = array_ops.split(value=concat, num_or_size_splits=4, axis=1)

            if self._use_peepholes:
                w_f_diag = vs.get_variable("w_f_diag", shape=[num_units], dtype=dtype)
                w_i_diag = vs.get_variable("w_i_diag", shape=[num_units], dtype=dtype)
                w_o_diag = vs.get_variable("w_o_diag", shape=[num_units], dtype=dtype)

                new_c = (
                        c * sigmoid(f + self._forget_bias + w_f_diag * c) +
                        sigmoid(i + w_i_diag * c) * self._activation(j))
            else:
                new_c = (
                        c * sigmoid(f + self._forget_bias) +
                        sigmoid(i) * self._activation(j))

            if self._cell_clip is not None:
                # pylint: disable=invalid-unary-operand-type
                new_c = clip_ops.clip_by_value(new_c, -self._cell_clip, self._cell_clip)
                # pylint: enable=invalid-unary-operand-type
            if self._use_peepholes:
                new_h = sigmoid(o + w_o_diag * new_c) * self._activation(new_c)
            else:
                new_h = sigmoid(o) * self._activation(new_c)

            if self._num_proj is not None:
                with vs.variable_scope("projection"):
                    new_h = self._linear(
                        new_h, self._num_proj, norm=self._norm, bias=False)

                if self._proj_clip is not None:
                    # pylint: disable=invalid-unary-operand-type
                    new_h = clip_ops.clip_by_value(new_h, -self._proj_clip,
                                                   self._proj_clip)
                    # pylint: enable=invalid-unary-operand-type

            new_state = rnn_cell_impl.LSTMStateTuple(new_c, new_h)
            return new_h, new_state


class IndRNNCell(rnn_cell_impl.LayerRNNCell):
    """Independently Recurrent Neural Network (IndRNN) cell
      (cf. https://arxiv.org/abs/1803.04831).

    Args:
      num_units: int, The number of units in the RNN cell.
      activation: Nonlinearity to use.  Default: `tanh`.
      reuse: (optional) Python boolean describing whether to reuse variables
       in an existing scope.  If not `True`, and the existing scope already has
       the given variables, an error is raised.
      name: String, the name of the layer. Layers with the same name will
        share weights, but to avoid mistakes we require reuse=True in such
        cases.
      dtype: Default dtype of the layer (default of `None` means use the type
        of the first input). Required when `build` is called before `call`.
    """

    def __init__(self,
                 num_units,
                 activation=None,
                 reuse=None,
                 name=None,
                 dtype=None):
        super(IndRNNCell, self).__init__(_reuse=reuse, name=name, dtype=dtype)

        # Inputs must be 2-dimensional.
        self.input_spec = input_spec.InputSpec(ndim=2)

        self._num_units = num_units
        self._activation = activation or math_ops.tanh

    @property
    def state_size(self):
        return self._num_units

    @property
    def output_size(self):
        return self._num_units

    def build(self, inputs_shape):
        if tensor_shape.dimension_value(inputs_shape[1]) is None:
            raise ValueError(
                "Expected inputs.shape[-1] to be known, saw shape: %s" % inputs_shape)

        input_depth = tensor_shape.dimension_value(inputs_shape[1])
        # pylint: disable=protected-access
        self._kernel_w = self.add_variable(
            "%s_w" % rnn_cell_impl._WEIGHTS_VARIABLE_NAME,
            shape=[input_depth, self._num_units])
        self._kernel_u = self.add_variable(
            "%s_u" % rnn_cell_impl._WEIGHTS_VARIABLE_NAME,
            shape=[1, self._num_units],
            initializer=init_ops.random_uniform_initializer(
                minval=-1, maxval=1, dtype=self.dtype))
        self._bias = self.add_variable(
            rnn_cell_impl._BIAS_VARIABLE_NAME,
            shape=[self._num_units],
            initializer=init_ops.zeros_initializer(dtype=self.dtype))
        # pylint: enable=protected-access

        self.built = True

    def call(self, inputs, state):
        """IndRNN: output = new_state = act(W * input + u * state + B)."""

        gate_inputs = math_ops.matmul(inputs, self._kernel_w) + (
                state * self._kernel_u)
        gate_inputs = nn_ops.bias_add(gate_inputs, self._bias)
        output = self._activation(gate_inputs)
        return output, output


class IndyGRUCell(rnn_cell_impl.LayerRNNCell):
    r"""Independently Gated Recurrent Unit cell.

    Based on IndRNNs (https://arxiv.org/abs/1803.04831) and similar to GRUCell,
    yet with the \\(U_r\\), \\(U_z\\), and \\(U\\) matrices in equations 5, 6, and
    8 of http://arxiv.org/abs/1406.1078 respectively replaced by diagonal
    matrices, i.e. a Hadamard product with a single vector:

      $$r_j = \sigma\left([\mathbf W_r\mathbf x]_j +
        [\mathbf u_r\circ \mathbf h_{(t-1)}]_j\right)$$
      $$z_j = \sigma\left([\mathbf W_z\mathbf x]_j +
        [\mathbf u_z\circ \mathbf h_{(t-1)}]_j\right)$$
      $$\tilde{h}^{(t)}_j = \phi\left([\mathbf W \mathbf x]_j +
        [\mathbf u \circ \mathbf r \circ \mathbf h_{(t-1)}]_j\right)$$

    where \\(\circ\\) denotes the Hadamard operator. This means that each IndyGRU
    node sees only its own state, as opposed to seeing all states in the same
    layer.

    Args:
      num_units: int, The number of units in the GRU cell.
      activation: Nonlinearity to use.  Default: `tanh`.
      reuse: (optional) Python boolean describing whether to reuse variables
       in an existing scope.  If not `True`, and the existing scope already has
       the given variables, an error is raised.
      kernel_initializer: (optional) The initializer to use for the weight
        matrices applied to the input.
      bias_initializer: (optional) The initializer to use for the bias.
      name: String, the name of the layer. Layers with the same name will
        share weights, but to avoid mistakes we require reuse=True in such
        cases.
      dtype: Default dtype of the layer (default of `None` means use the type
        of the first input). Required when `build` is called before `call`.
    """

    def __init__(self,
                 num_units,
                 activation=None,
                 reuse=None,
                 kernel_initializer=None,
                 bias_initializer=None,
                 name=None,
                 dtype=None):
        super(IndyGRUCell, self).__init__(_reuse=reuse, name=name, dtype=dtype)

        # Inputs must be 2-dimensional.
        self.input_spec = input_spec.InputSpec(ndim=2)

        self._num_units = num_units
        self._activation = activation or math_ops.tanh
        self._kernel_initializer = kernel_initializer
        self._bias_initializer = bias_initializer

    @property
    def state_size(self):
        return self._num_units

    @property
    def output_size(self):
        return self._num_units

    def build(self, inputs_shape):
        if tensor_shape.dimension_value(inputs_shape[1]) is None:
            raise ValueError(
                "Expected inputs.shape[-1] to be known, saw shape: %s" % inputs_shape)

        input_depth = tensor_shape.dimension_value(inputs_shape[1])
        # pylint: disable=protected-access
        self._gate_kernel_w = self.add_variable(
            "gates/%s_w" % rnn_cell_impl._WEIGHTS_VARIABLE_NAME,
            shape=[input_depth, 2 * self._num_units],
            initializer=self._kernel_initializer)
        self._gate_kernel_u = self.add_variable(
            "gates/%s_u" % rnn_cell_impl._WEIGHTS_VARIABLE_NAME,
            shape=[1, 2 * self._num_units],
            initializer=init_ops.random_uniform_initializer(
                minval=-1, maxval=1, dtype=self.dtype))
        self._gate_bias = self.add_variable(
            "gates/%s" % rnn_cell_impl._BIAS_VARIABLE_NAME,
            shape=[2 * self._num_units],
            initializer=(self._bias_initializer
                         if self._bias_initializer is not None else
                         init_ops.constant_initializer(1.0, dtype=self.dtype)))
        self._candidate_kernel_w = self.add_variable(
            "candidate/%s" % rnn_cell_impl._WEIGHTS_VARIABLE_NAME,
            shape=[input_depth, self._num_units],
            initializer=self._kernel_initializer)
        self._candidate_kernel_u = self.add_variable(
            "candidate/%s_u" % rnn_cell_impl._WEIGHTS_VARIABLE_NAME,
            shape=[1, self._num_units],
            initializer=init_ops.random_uniform_initializer(
                minval=-1, maxval=1, dtype=self.dtype))
        self._candidate_bias = self.add_variable(
            "candidate/%s" % rnn_cell_impl._BIAS_VARIABLE_NAME,
            shape=[self._num_units],
            initializer=(self._bias_initializer
                         if self._bias_initializer is not None else
                         init_ops.zeros_initializer(dtype=self.dtype)))
        # pylint: enable=protected-access

        self.built = True

    def call(self, inputs, state):
        """Recurrently independent Gated Recurrent Unit (GRU) with nunits cells."""

        gate_inputs = math_ops.matmul(inputs, self._gate_kernel_w) + (
                gen_array_ops.tile(state, [1, 2]) * self._gate_kernel_u)
        gate_inputs = nn_ops.bias_add(gate_inputs, self._gate_bias)

        value = math_ops.sigmoid(gate_inputs)
        r, u = array_ops.split(value=value, num_or_size_splits=2, axis=1)

        r_state = r * state

        candidate = math_ops.matmul(inputs, self._candidate_kernel_w) + (
                r_state * self._candidate_kernel_u)
        candidate = nn_ops.bias_add(candidate, self._candidate_bias)

        c = self._activation(candidate)
        new_h = u * state + (1 - u) * c
        return new_h, new_h


class IndyLSTMCell(rnn_cell_impl.LayerRNNCell):
    r"""Basic IndyLSTM recurrent network cell.

    Based on IndRNNs (https://arxiv.org/abs/1803.04831) and similar to
    BasicLSTMCell, yet with the \\(U_f\\), \\(U_i\\), \\(U_o\\) and \\(U_c\\)
    matrices in the regular LSTM equations replaced by diagonal matrices, i.e. a
    Hadamard product with a single vector:

      $$f_t = \sigma_g\left(W_f x_t + u_f \circ h_{t-1} + b_f\right)$$
      $$i_t = \sigma_g\left(W_i x_t + u_i \circ h_{t-1} + b_i\right)$$
      $$o_t = \sigma_g\left(W_o x_t + u_o \circ h_{t-1} + b_o\right)$$
      $$c_t = f_t \circ c_{t-1} +
              i_t \circ \sigma_c\left(W_c x_t + u_c \circ h_{t-1} + b_c\right)$$

    where \\(\circ\\) denotes the Hadamard operator. This means that each IndyLSTM
    node sees only its own state \\(h\\) and \\(c\\), as opposed to seeing all
    states in the same layer.

    We add forget_bias (default: 1) to the biases of the forget gate in order to
    reduce the scale of forgetting in the beginning of the training.

    It does not allow cell clipping, a projection layer, and does not
    use peep-hole connections: it is the basic baseline.

    For a detailed analysis of IndyLSTMs, see https://arxiv.org/abs/1903.08023.
    """

    def __init__(self,
                 num_units,
                 forget_bias=1.0,
                 activation=None,
                 reuse=None,
                 kernel_initializer=None,
                 bias_initializer=None,
                 name=None,
                 dtype=None):
        """Initialize the IndyLSTM cell.

        Args:
          num_units: int, The number of units in the LSTM cell.
          forget_bias: float, The bias added to forget gates (see above).
            Must set to `0.0` manually when restoring from CudnnLSTM-trained
            checkpoints.
          activation: Activation function of the inner states.  Default: `tanh`.
          reuse: (optional) Python boolean describing whether to reuse variables
            in an existing scope.  If not `True`, and the existing scope already has
            the given variables, an error is raised.
          kernel_initializer: (optional) The initializer to use for the weight
            matrix applied to the inputs.
          bias_initializer: (optional) The initializer to use for the bias.
          name: String, the name of the layer. Layers with the same name will
            share weights, but to avoid mistakes we require reuse=True in such
            cases.
          dtype: Default dtype of the layer (default of `None` means use the type
            of the first input). Required when `build` is called before `call`.
        """
        super(IndyLSTMCell, self).__init__(_reuse=reuse, name=name, dtype=dtype)

        # Inputs must be 2-dimensional.
        self.input_spec = input_spec.InputSpec(ndim=2)

        self._num_units = num_units
        self._forget_bias = forget_bias
        self._activation = activation or math_ops.tanh
        self._kernel_initializer = kernel_initializer
        self._bias_initializer = bias_initializer

    @property
    def state_size(self):
        return rnn_cell_impl.LSTMStateTuple(self._num_units, self._num_units)

    @property
    def output_size(self):
        return self._num_units

    def build(self, inputs_shape):
        if tensor_shape.dimension_value(inputs_shape[1]) is None:
            raise ValueError(
                "Expected inputs.shape[-1] to be known, saw shape: %s" % inputs_shape)

        input_depth = tensor_shape.dimension_value(inputs_shape[1])
        # pylint: disable=protected-access
        self._kernel_w = self.add_variable(
            "%s_w" % rnn_cell_impl._WEIGHTS_VARIABLE_NAME,
            shape=[input_depth, 4 * self._num_units],
            initializer=self._kernel_initializer)
        self._kernel_u = self.add_variable(
            "%s_u" % rnn_cell_impl._WEIGHTS_VARIABLE_NAME,
            shape=[1, 4 * self._num_units],
            initializer=init_ops.random_uniform_initializer(
                minval=-1, maxval=1, dtype=self.dtype))
        self._bias = self.add_variable(
            rnn_cell_impl._BIAS_VARIABLE_NAME,
            shape=[4 * self._num_units],
            initializer=(self._bias_initializer
                         if self._bias_initializer is not None else
                         init_ops.zeros_initializer(dtype=self.dtype)))
        # pylint: enable=protected-access

        self.built = True

    def call(self, inputs, state):
        """Independent Long short-term memory cell (IndyLSTM).

        Args:
          inputs: `2-D` tensor with shape `[batch_size, input_size]`.
          state: An `LSTMStateTuple` of state tensors, each shaped
            `[batch_size, num_units]`.

        Returns:
          A pair containing the new hidden state, and the new state (a
            `LSTMStateTuple`).
        """
        sigmoid = math_ops.sigmoid
        one = constant_op.constant(1, dtype=dtypes.int32)
        c, h = state

        gate_inputs = math_ops.matmul(inputs, self._kernel_w)
        gate_inputs += gen_array_ops.tile(h, [1, 4]) * self._kernel_u
        gate_inputs = nn_ops.bias_add(gate_inputs, self._bias)

        # i = input_gate, j = new_input, f = forget_gate, o = output_gate
        i, j, f, o = array_ops.split(
            value=gate_inputs, num_or_size_splits=4, axis=one)

        forget_bias_tensor = constant_op.constant(self._forget_bias, dtype=f.dtype)
        # Note that using `add` and `multiply` instead of `+` and `*` gives a
        # performance improvement. So using those at the cost of readability.
        add = math_ops.add
        multiply = math_ops.multiply
        new_c = add(
            multiply(c, sigmoid(add(f, forget_bias_tensor))),
            multiply(sigmoid(i), self._activation(j)))
        new_h = multiply(self._activation(new_c), sigmoid(o))

        new_state = rnn_cell_impl.LSTMStateTuple(new_c, new_h)
        return new_h, new_state


NTMControllerState = collections.namedtuple(
    "NTMControllerState",
    ("controller_state", "read_vector_list", "w_list", "M", "time"))


class NTMCell(rnn_cell_impl.LayerRNNCell):
    """Neural Turing Machine Cell with RNN controller.

      Implementation based on:
      https://arxiv.org/abs/1807.08518
      Mark Collier, Joeran Beel

      which is in turn based on the source code of:
      https://github.com/snowkylin/ntm

      and of course the original NTM paper:
      Neural Turing Machines
      https://arxiv.org/abs/1410.5401
      A Graves, G Wayne, I Danihelka
    """

    def __init__(self,
                 controller,
                 memory_size,
                 memory_vector_dim,
                 read_head_num,
                 write_head_num,
                 shift_range=1,
                 output_dim=None,
                 clip_value=20,
                 dtype=dtypes.float32,
                 name=None):
        """Initialize the NTM Cell.

          Args:
            controller: an RNNCell, the RNN controller.
            memory_size: int, The number of memory locations in the NTM memory
              matrix
            memory_vector_dim: int, The dimensionality of each location in the NTM
              memory matrix
            read_head_num: int, The number of read heads from the controller into
              memory
            write_head_num: int, The number of write heads from the controller into
              memory
            shift_range: int, The number of places to the left/right it is possible
              to iterate the previous address to in a single step
            output_dim: int, The number of dimensions to make a linear projection of
              the NTM controller outputs to. If None, no linear projection is
              applied
            clip_value: float, The maximum absolute value the controller parameters
              are clipped to
            dtype: Default dtype of the layer (default of `None` means use the type
              of the first input). Required when `build` is called before `call`.
            name: String, the name of the layer. Layers with the same name will
              share weights, but to avoid mistakes we require reuse=True in such
              cases.
        """
        super(NTMCell, self).__init__(dtype=dtype, name=name)

        rnn_cell_impl.assert_like_rnncell("NTM RNN controller cell", controller)

        self.controller = controller
        self.memory_size = memory_size
        self.memory_vector_dim = memory_vector_dim
        self.read_head_num = read_head_num
        self.write_head_num = write_head_num
        self.clip_value = clip_value

        self.output_dim = output_dim
        self.shift_range = shift_range

        self.num_parameters_per_head = (
                self.memory_vector_dim + 2 * self.shift_range + 4)
        self.num_heads = self.read_head_num + self.write_head_num
        self.total_parameter_num = (
                self.num_parameters_per_head * self.num_heads +
                self.memory_vector_dim * 2 * self.write_head_num)

    @property
    def state_size(self):
        return NTMControllerState(
            controller_state=self.controller.state_size,
            read_vector_list=[
                self.memory_vector_dim for _ in range(self.read_head_num)
            ],
            w_list=[
                self.memory_size
                for _ in range(self.read_head_num + self.write_head_num)
            ],
            M=tensor_shape.TensorShape([self.memory_size * self.memory_vector_dim]),
            time=tensor_shape.TensorShape([]))

    @property
    def output_size(self):
        return self.output_dim

    def build(self, inputs_shape):
        if self.output_dim is None:
            if inputs_shape[1].value is None:
                raise ValueError(
                    "Expected inputs.shape[-1] to be known, saw shape: %s" %
                    inputs_shape)
            else:
                self.output_dim = inputs_shape[1].value

        def _create_linear_initializer(input_size, dtype=dtypes.float32):
            stddev = 1.0 / math.sqrt(input_size)
            return init_ops.truncated_normal_initializer(stddev=stddev, dtype=dtype)

        self._params_kernel = self.add_variable(
            "parameters_kernel",
            shape=[self.controller.output_size, self.total_parameter_num],
            initializer=_create_linear_initializer(self.controller.output_size))

        self._params_bias = self.add_variable(
            "parameters_bias",
            shape=[self.total_parameter_num],
            initializer=init_ops.constant_initializer(0.0, dtype=self.dtype))

        self._output_kernel = self.add_variable(
            "output_kernel",
            shape=[
                self.controller.output_size +
                self.memory_vector_dim * self.read_head_num, self.output_dim
            ],
            initializer=_create_linear_initializer(self.controller.output_size +
                                                   self.memory_vector_dim *
                                                   self.read_head_num))

        self._output_bias = self.add_variable(
            "output_bias",
            shape=[self.output_dim],
            initializer=init_ops.constant_initializer(0.0, dtype=self.dtype))

        self._init_read_vectors = [
            self.add_variable(
                "initial_read_vector_%d" % i,
                shape=[1, self.memory_vector_dim],
                initializer=initializers.glorot_uniform())
            for i in range(self.read_head_num)
        ]

        self._init_address_weights = [
            self.add_variable(
                "initial_address_weights_%d" % i,
                shape=[1, self.memory_size],
                initializer=initializers.glorot_uniform())
            for i in range(self.read_head_num + self.write_head_num)
        ]

        self._M = self.add_variable(
            "memory",
            shape=[self.memory_size, self.memory_vector_dim],
            initializer=init_ops.constant_initializer(1e-6, dtype=self.dtype))

        self.built = True

    def call(self, x, prev_state):
        # Addressing Mechanisms (Sec 3.3)

        def _prev_read_vector_list_initial_value():
            return [
                self._expand(
                    math_ops.tanh(
                        array_ops.squeeze(
                            math_ops.matmul(
                                array_ops.ones([1, 1]), self._init_read_vectors[i]))),
                    dim=0,
                    N=x.shape[0].value or array_ops.shape(x)[0])
                for i in range(self.read_head_num)
            ]

        prev_read_vector_list = control_flow_ops.cond(
            math_ops.equal(prev_state.time,
                           0), _prev_read_vector_list_initial_value, lambda:
            prev_state.read_vector_list)
        if self.read_head_num == 1:
            prev_read_vector_list = [prev_read_vector_list]

        controller_input = array_ops.concat([x] + prev_read_vector_list, axis=1)
        controller_output, controller_state = self.controller(
            controller_input, prev_state.controller_state)

        parameters = math_ops.matmul(controller_output, self._params_kernel)
        parameters = nn_ops.bias_add(parameters, self._params_bias)
        parameters = clip_ops.clip_by_value(parameters, -self.clip_value,
                                            self.clip_value)
        head_parameter_list = array_ops.split(
            parameters[:, :self.num_parameters_per_head * self.num_heads],
            self.num_heads,
            axis=1)
        erase_add_list = array_ops.split(
            parameters[:, self.num_parameters_per_head * self.num_heads:],
            2 * self.write_head_num,
            axis=1)

        def _prev_w_list_initial_value():
            return [
                self._expand(
                    nn_ops.softmax(
                        array_ops.squeeze(
                            math_ops.matmul(
                                array_ops.ones([1, 1]),
                                self._init_address_weights[i]))),
                    dim=0,
                    N=x.shape[0].value or array_ops.shape(x)[0])
                for i in range(self.read_head_num + self.write_head_num)
            ]

        prev_w_list = control_flow_ops.cond(
            math_ops.equal(prev_state.time, 0),
            _prev_w_list_initial_value, lambda: prev_state.w_list)
        if (self.read_head_num + self.write_head_num) == 1:
            prev_w_list = [prev_w_list]

        prev_M = control_flow_ops.cond(
            math_ops.equal(prev_state.time, 0), lambda: self._expand(
                self._M, dim=0, N=x.shape[0].value or array_ops.shape(x)[0]),
            lambda: prev_state.M)

        w_list = []
        for i, head_parameter in enumerate(head_parameter_list):
            k = math_ops.tanh(head_parameter[:, 0:self.memory_vector_dim])
            beta = nn_ops.softplus(head_parameter[:, self.memory_vector_dim])
            g = math_ops.sigmoid(head_parameter[:, self.memory_vector_dim + 1])
            s = nn_ops.softmax(head_parameter[:, self.memory_vector_dim +
                                                 2:(self.memory_vector_dim + 2 +
                                                    (self.shift_range * 2 + 1))])
            gamma = nn_ops.softplus(head_parameter[:, -1]) + 1
            w = self._addressing(k, beta, g, s, gamma, prev_M, prev_w_list[i])
            w_list.append(w)

        # Reading (Sec 3.1)

        read_w_list = w_list[:self.read_head_num]
        read_vector_list = []
        for i in range(self.read_head_num):
            read_vector = math_ops.reduce_sum(
                array_ops.expand_dims(read_w_list[i], dim=2) * prev_M, axis=1)
            read_vector_list.append(read_vector)

        # Writing (Sec 3.2)

        write_w_list = w_list[self.read_head_num:]
        M = prev_M
        for i in range(self.write_head_num):
            w = array_ops.expand_dims(write_w_list[i], axis=2)
            erase_vector = array_ops.expand_dims(
                math_ops.sigmoid(erase_add_list[i * 2]), axis=1)
            add_vector = array_ops.expand_dims(
                math_ops.tanh(erase_add_list[i * 2 + 1]), axis=1)
            erase_M = array_ops.ones_like(M) - math_ops.matmul(w, erase_vector)
            M = M * erase_M + math_ops.matmul(w, add_vector)

        output = math_ops.matmul(
            array_ops.concat([controller_output] + read_vector_list, axis=1),
            self._output_kernel)
        output = nn_ops.bias_add(output, self._output_bias)
        output = clip_ops.clip_by_value(output, -self.clip_value, self.clip_value)

        return output, NTMControllerState(
            controller_state=controller_state,
            read_vector_list=read_vector_list,
            w_list=w_list,
            M=M,
            time=prev_state.time + 1)

    def _expand(self, x, dim, N):
        return array_ops.concat([array_ops.expand_dims(x, dim) for _ in range(N)],
                                axis=dim)

    def _addressing(self, k, beta, g, s, gamma, prev_M, prev_w):
        # Sec 3.3.1 Focusing by Content

        k = array_ops.expand_dims(k, axis=2)
        inner_product = math_ops.matmul(prev_M, k)
        k_norm = math_ops.sqrt(
            math_ops.reduce_sum(math_ops.square(k), axis=1, keepdims=True))
        M_norm = math_ops.sqrt(
            math_ops.reduce_sum(math_ops.square(prev_M), axis=2, keepdims=True))
        norm_product = M_norm * k_norm

        # eq (6)
        K = array_ops.squeeze(inner_product / (norm_product + 1e-8))

        K_amplified = math_ops.exp(array_ops.expand_dims(beta, axis=1) * K)

        # eq (5)
        w_c = K_amplified / math_ops.reduce_sum(K_amplified, axis=1, keepdims=True)

        # Sec 3.3.2 Focusing by Location

        g = array_ops.expand_dims(g, axis=1)

        # eq (7)
        w_g = g * w_c + (1 - g) * prev_w

        s = array_ops.concat([
            s[:, :self.shift_range + 1],
            array_ops.zeros([
                s.shape[0].value or array_ops.shape(s)[0], self.memory_size -
                (self.shift_range * 2 + 1)
            ]), s[:, -self.shift_range:]
        ],
            axis=1)
        t = array_ops.concat(
            [array_ops.reverse(s, axis=[1]),
             array_ops.reverse(s, axis=[1])],
            axis=1)
        s_matrix = array_ops.stack([
            t[:, self.memory_size - i - 1:self.memory_size * 2 - i - 1]
            for i in range(self.memory_size)
        ],
            axis=1)

        # eq (8)
        w_ = math_ops.reduce_sum(
            array_ops.expand_dims(w_g, axis=1) * s_matrix, axis=2)
        w_sharpen = math_ops.pow(w_, array_ops.expand_dims(gamma, axis=1))

        # eq (9)
        w = w_sharpen / math_ops.reduce_sum(w_sharpen, axis=1, keepdims=True)

        return w

    def zero_state(self, batch_size, dtype):
        read_vector_list = [
            array_ops.zeros([batch_size, self.memory_vector_dim])
            for _ in range(self.read_head_num)
        ]

        w_list = [
            array_ops.zeros([batch_size, self.memory_size])
            for _ in range(self.read_head_num + self.write_head_num)
        ]

        controller_init_state = self.controller.zero_state(batch_size, dtype)

        M = array_ops.zeros([batch_size, self.memory_size, self.memory_vector_dim])

        return NTMControllerState(
            controller_state=controller_init_state,
            read_vector_list=read_vector_list,
            w_list=w_list,
            M=M,
            time=0)


class MinimalRNNCell(rnn_cell_impl.LayerRNNCell):
    """MinimalRNN cell.

    The implementation is based on:

      https://arxiv.org/pdf/1806.05394v2.pdf

    Minmin Chen, Jeffrey Pennington, Samuel S. Schoenholz.
    "Dynamical Isometry and a Mean Field Theory of RNNs: Gating Enables Signal
     Propagation in Recurrent Neural Networks." ICML, 2018.

    A MinimalRNN cell first projects the input to the hidden space. The new
    hidden state is then calculated as a weighted sum of the projected input and
    the previous hidden state, using a single update gate.
    """

    def __init__(self,
                 units,
                 activation="tanh",
                 kernel_initializer="glorot_uniform",
                 bias_initializer="ones",
                 name=None,
                 dtype=None,
                 **kwargs):
        """Initialize the parameters for a MinimalRNN cell.

        Args:
          units: int, The number of units in the MinimalRNN cell.
          activation: Nonlinearity to use in the feedforward network. Default:
            `tanh`.
          kernel_initializer: The initializer to use for the weight in the update
            gate and feedforward network. Default: `glorot_uniform`.
          bias_initializer: The initializer to use for the bias in the update
            gate. Default: `ones`.
          name: String, the name of the cell.
          dtype: Default dtype of the cell.
          **kwargs: Dict, keyword named properties for common cell attributes.
        """
        super(MinimalRNNCell, self).__init__(name=name, dtype=dtype, **kwargs)

        # Inputs must be 2-dimensional.
        self.input_spec = input_spec.InputSpec(ndim=2)

        self.units = units
        self.activation = activations.get(activation)
        self.kernel_initializer = initializers.get(kernel_initializer)
        self.bias_initializer = initializers.get(bias_initializer)

    @property
    def state_size(self):
        return self.units

    @property
    def output_size(self):
        return self.units

    def build(self, inputs_shape):
        if inputs_shape[-1] is None:
            raise ValueError("Expected inputs.shape[-1] to be known, saw shape: %s"
                             % str(inputs_shape))

        input_size = inputs_shape[-1]
        # pylint: disable=protected-access
        # self._kernel contains W_x, W, V
        self.kernel = self.add_weight(
            name=rnn_cell_impl._WEIGHTS_VARIABLE_NAME,
            shape=[input_size + 2 * self.units, self.units],
            initializer=self.kernel_initializer)
        self.bias = self.add_weight(
            name=rnn_cell_impl._BIAS_VARIABLE_NAME,
            shape=[self.units],
            initializer=self.bias_initializer)
        # pylint: enable=protected-access

        self.built = True

    def call(self, inputs, state):
        """Run one step of MinimalRNN.

        Args:
          inputs: input Tensor, must be 2-D, `[batch, input_size]`.
          state: state Tensor, must be 2-D, `[batch, state_size]`.

        Returns:
          A tuple containing:

          - Output: A `2-D` tensor with shape `[batch_size, state_size]`.
          - New state: A `2-D` tensor with shape `[batch_size, state_size]`.

        Raises:
          ValueError: If input size cannot be inferred from inputs via
            static shape inference.
        """
        input_size = inputs.get_shape()[1]
        if tensor_shape.dimension_value(input_size) is None:
            raise ValueError("Could not infer input size from inputs.get_shape()[-1]")

        feedforward_weight, gate_weight = array_ops.split(
            value=self.kernel,
            num_or_size_splits=[tensor_shape.dimension_value(input_size),
                                2 * self.units],
            axis=0)

        feedforward = math_ops.matmul(inputs, feedforward_weight)
        feedforward = self.activation(feedforward)

        gate_inputs = math_ops.matmul(
            array_ops.concat([feedforward, state], 1), gate_weight)
        gate_inputs = nn_ops.bias_add(gate_inputs, self.bias)
        u = math_ops.sigmoid(gate_inputs)

        new_h = u * state + (1 - u) * feedforward
        return new_h, new_h


class CFNCell(rnn_cell_impl.LayerRNNCell):
    """Chaos Free Network cell.

    The implementation is based on:

      https://openreview.net/pdf?id=S1dIzvclg

    Thomas Laurent, James von Brecht.
    "A recurrent neural network without chaos." ICLR, 2017.

    A CFN cell first projects the input to the hidden space. The hidden state
    goes through a contractive mapping. The new hidden state is then calculated
    as a linear combination of the projected input and the contracted previous
    hidden state, using decoupled input and forget gates.
    """

    def __init__(self,
                 units,
                 activation="tanh",
                 kernel_initializer="glorot_uniform",
                 bias_initializer="ones",
                 name=None,
                 dtype=None,
                 **kwargs):
        """Initialize the parameters for a CFN cell.

        Args:
          units: int, The number of units in the CFN cell.
          activation: Nonlinearity to use. Default: `tanh`.
          kernel_initializer: Initializer for the `kernel` weights
            matrix. Default: `glorot_uniform`.
          bias_initializer: The initializer to use for the bias in the
            gates. Default: `ones`.
          name: String, the name of the cell.
          dtype: Default dtype of the cell.
          **kwargs: Dict, keyword named properties for common cell attributes.
        """
        super(CFNCell, self).__init__(name=name, dtype=dtype, **kwargs)

        # Inputs must be 2-dimensional.
        self.input_spec = input_spec.InputSpec(ndim=2)

        self.units = units
        self.activation = activations.get(activation)
        self.kernel_initializer = initializers.get(kernel_initializer)
        self.bias_initializer = initializers.get(bias_initializer)

    @property
    def state_size(self):
        return self.units

    @property
    def output_size(self):
        return self.units

    def build(self, inputs_shape):
        if inputs_shape[-1] is None:
            raise ValueError("Expected inputs.shape[-1] to be known, saw shape: %s"
                             % str(inputs_shape))

        input_size = inputs_shape[-1]
        # pylint: disable=protected-access
        # `self.kernel` contains V_{\theta}, V_{\eta}, W.
        # `self.recurrent_kernel` contains U_{\theta}, U_{\eta}.
        # `self.bias` contains b_{\theta}, b_{\eta}.
        self.kernel = self.add_weight(
            shape=[input_size, 3 * self.units],
            name=rnn_cell_impl._WEIGHTS_VARIABLE_NAME,
            initializer=self.kernel_initializer)
        self.recurrent_kernel = self.add_weight(
            shape=[self.units, 2 * self.units],
            name="recurrent_%s" % rnn_cell_impl._WEIGHTS_VARIABLE_NAME,
            initializer=self.kernel_initializer)
        self.bias = self.add_weight(
            shape=[2 * self.units],
            name=rnn_cell_impl._BIAS_VARIABLE_NAME,
            initializer=self.bias_initializer)
        # pylint: enable=protected-access

        self.built = True

    def call(self, inputs, state):
        """Run one step of CFN.

        Args:
          inputs: input Tensor, must be 2-D, `[batch, input_size]`.
          state: state Tensor, must be 2-D, `[batch, state_size]`.

        Returns:
          A tuple containing:

          - Output: A `2-D` tensor with shape `[batch_size, state_size]`.
          - New state: A `2-D` tensor with shape `[batch_size, state_size]`.

        Raises:
          ValueError: If input size cannot be inferred from inputs via
            static shape inference.
        """
        input_size = inputs.get_shape()[-1]
        if tensor_shape.dimension_value(input_size) is None:
            raise ValueError("Could not infer input size from inputs.get_shape()[-1]")

        # The variable names u, v, w, b are consistent with the notations in the
        # original paper.
        v, w = array_ops.split(
            value=self.kernel,
            num_or_size_splits=[2 * self.units, self.units],
            axis=1)
        u = self.recurrent_kernel
        b = self.bias

        gates = math_ops.matmul(state, u) + math_ops.matmul(inputs, v)
        gates = nn_ops.bias_add(gates, b)
        gates = math_ops.sigmoid(gates)
        theta, eta = array_ops.split(value=gates,
                                     num_or_size_splits=2,
                                     axis=1)

        proj_input = math_ops.matmul(inputs, w)

        # The input gate is (1 - eta), which is different from the original paper.
        # This is for the propose of initialization. With the default
        # bias_initializer `ones`, the input gate is initialized to a small number.
        new_h = theta * self.activation(state) + (1 - eta) * self.activation(
            proj_input)

        return new_h, new_h
