def __call__(self, inputs, state, scope=None):
        """Long short-term memory cell (LSTM)."""
        with tf.variable_scope(self, scope or "basic_lstm_cell", reuse=self._reuse):
            # Parameters of gates are concatenated into one multiply for
            # efficiency.
            if self._state_is_tuple:
                c_prev, h_prev = state
            else:
                c_prev, h_prev = tf.split(
                    value=state, num_or_size_splits=2, axis=1)
            concat = tf.contrib.rnn._linear(
                [inputs, h_prev], 4 * self._num_units, True)

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

            c = (c_prev * tf.sigmoid(f + self._forget_bias) +
                 tf.sigmoid(i) * tf.tanh(g))
            h = tf.tanh(c) * tf.sigmoid(o)

            if self._state_is_tuple:
                new_state = LSTMStateTuple(c, h)
            else:
                new_state = tf.concat([c, h], 1)
            return h, new_state 

def rnn_placeholders(state):
    """
    Given nested [multilayer] RNN state tensor, infers and returns state placeholders.

    Args:
        state:  tf.nn.lstm zero-state tuple.

    Returns:    tuple of placeholders
    """
    if isinstance(state, tf.contrib.rnn.LSTMStateTuple):
        c, h = state
        c = tf.placeholder(tf.float32, tf.TensorShape([None]).concatenate(c.get_shape()[1:]), c.op.name + '_c_pl')
        h = tf.placeholder(tf.float32, tf.TensorShape([None]).concatenate(h.get_shape()[1:]), h.op.name + '_h_pl')
        return tf.contrib.rnn.LSTMStateTuple(c, h)
    elif isinstance(state, tf.Tensor):
        h = state
        h = tf.placeholder(tf.float32, tf.TensorShape([None]).concatenate(h.get_shape()[1:]), h.op.name + '_h_pl')
        return h
    else:
        structure = [rnn_placeholders(x) for x in state]
        return tuple(structure) 

def _show_struct(struct):
    # Debug utility
    if isinstance(struct, dict):
        for key, value in struct.items():
            print(key)
            _show_struct(value)

    elif type(struct) in [LSTMStateTuple, tuple, list]:
        print('LSTM/tuple/list:', type(struct), len(struct))
        for i in struct:
            _show_struct(i)

    else:
        try:
            print('shape: {}, type: {}'.format(np.asarray(struct).shape, type(struct)))

        except AttributeError:
            print('value:', struct) 

def as_array(self, struct, squeeze_axis=None):
        if isinstance(struct, dict):
            out = {}
            for key, value in struct.items():
                out[key] = self.as_array(value, squeeze_axis)
            return out

        elif isinstance(struct, tuple):
            return tuple([self.as_array(value, squeeze_axis) for value in struct])

        elif isinstance(struct, LSTMStateTuple):
            return LSTMStateTuple(self.as_array(struct[0], squeeze_axis), self.as_array(struct[1], squeeze_axis))

        else:
            if squeeze_axis is not None:
                return np.squeeze(np.asarray(struct), axis=squeeze_axis)

            else:
                return np.asarray(struct) 

def mask_finished(finished, now_, prev_):
        mask = tf.expand_dims(tf.to_float(finished), 1)

        if isinstance(prev_, tuple):
            # tuple states
            next_ = []
            for ns, s in zip(now_, prev_):
                # fucking LSTMStateTuple
                if isinstance(ns, LSTMStateTuple):
                    next_.append(
                        LSTMStateTuple(c=(1. - mask) * ns.c + mask * s.c,
                                       h=(1. - mask) * ns.h + mask * s.h))
                else:
                    next_.append((1. - mask) * ns + mask * s)
            next_ = tuple(next_)
        else:
            next_ = (1. - mask) * now_ + mask * prev_

        return next_ 

def predict(self, batch, states, batchsize=1):
        """

        :param batch: batch of data
        :param states: prior states
        :param batchsize: batchsize if getting initial state
        :return: 
            predicted_outputs: softmax predictions
            states: states of RNN
        """

        if states is None:
            return self.session.run([self.predicted_outputs, self.rnn_states], {self.inputs: batch})
        else:
            c = states[0]
            h = states[1]
            return self.session.run([self.predicted_outputs, self.rnn_states],
                                    {self.inputs: batch,
                                     self.initial_state: LSTMStateTuple(c, h)}) 

def build_lstm(x, size, name, step_size):
    lstm = rnn.BasicLSTMCell(size, state_is_tuple=True)

    c_init = np.zeros((1, lstm.state_size.c), np.float32)
    h_init = np.zeros((1, lstm.state_size.h), np.float32)
    state_init = [c_init, h_init]

    c_in = tf.placeholder(tf.float32, 
            shape=[1, lstm.state_size.c],
            name='c_in')
    h_in = tf.placeholder(tf.float32, 
            shape=[1, lstm.state_size.h],
            name='h_in')
    state_in = [c_in, h_in]

    state_in = rnn.LSTMStateTuple(c_in, h_in)

    lstm_outputs, lstm_state = tf.nn.dynamic_rnn(
        lstm, x, initial_state=state_in, sequence_length=step_size,
        time_major=False)
    lstm_outputs = tf.reshape(lstm_outputs, [-1, size])

    lstm_c, lstm_h = lstm_state
    state_out = [lstm_c[:1, :], lstm_h[:1, :]]
    return lstm_outputs, state_init, state_in, state_out 

def __init__(self,x,size,step_size):
        lstm = rnn.BasicLSTMCell(size, state_is_tuple=True)

        c_init = np.zeros((1, lstm.state_size.c), np.float32)
        h_init = np.zeros((1, lstm.state_size.h), np.float32)
        self.state_init = [c_init, h_init]

        c_in = tf.placeholder(tf.float32, 
                shape=[1, lstm.state_size.c],
                name='c_in')
        h_in = tf.placeholder(tf.float32, 
                shape=[1, lstm.state_size.h],
                name='h_in')
        self.state_in = [c_in, h_in]

        state_in = rnn.LSTMStateTuple(c_in, h_in)

        lstm_outputs, lstm_state = tf.nn.dynamic_rnn(
            lstm, x, initial_state=state_in, sequence_length=step_size,
            time_major=False)
        lstm_outputs = tf.reshape(lstm_outputs, [-1, size])

        lstm_c, lstm_h = lstm_state
        self.state_out = [lstm_c[:1, :], lstm_h[:1, :]]
        self.output = lstm_outputs 

def __call__(self, inputs, state, scope=None):
        with tf.variable_scope(self.scope):
            c, h = state
            h = dropout(h, self.keep_recurrent_probs, self.is_train)

            mat = _compute_gates(inputs, h, self.num_units, self.forget_bias,
                                        self.kernel_initializer, self.recurrent_initializer, True)

            i, j, f, o = tf.split(value=mat, num_or_size_splits=4, axis=1)

            new_c = (c * self.recurrent_activation(f) + self.recurrent_activation(i) *
                     self.activation(j))
            new_h = self.activation(new_c) * self.recurrent_activation(o)

            new_state = LSTMStateTuple(new_c, new_h)

        return new_h, new_state 

def decoder(x, decoder_inputs, keep_prob, sequence_length, memory, memory_length, first_attention):
    with tf.variable_scope("Decoder") as scope:
        label_embeddings = tf.get_variable(name="embeddings", shape=[n_classes, embedding_size], dtype=tf.float32)
        train_inputs_embedded = tf.nn.embedding_lookup(label_embeddings, decoder_inputs)
        lstm = rnn.LayerNormBasicLSTMCell(n_hidden, dropout_keep_prob=keep_prob)
        output_l = layers_core.Dense(n_classes, use_bias=True)
        encoder_state = rnn.LSTMStateTuple(x, x)
        attention_mechanism = BahdanauAttention(embedding_size, memory=memory, memory_sequence_length=memory_length)
        cell = AttentionWrapper(lstm, attention_mechanism, output_attention=False)
        cell_state = cell.zero_state(dtype=tf.float32, batch_size=train_batch_size)
        cell_state = cell_state.clone(cell_state=encoder_state, attention=first_attention)
        train_helper = TrainingHelper(train_inputs_embedded, sequence_length)
        train_decoder = BasicDecoder(cell, train_helper, cell_state, output_layer=output_l)
        decoder_outputs_train, decoder_state_train, decoder_seq_train = dynamic_decode(train_decoder, impute_finished=True)
        tiled_inputs = tile_batch(memory, multiplier=beam_width)
        tiled_sequence_length = tile_batch(memory_length, multiplier=beam_width)
        tiled_first_attention = tile_batch(first_attention, multiplier=beam_width)
        attention_mechanism = BahdanauAttention(embedding_size, memory=tiled_inputs, memory_sequence_length=tiled_sequence_length)
        x2 = tile_batch(x, beam_width)
        encoder_state2 = rnn.LSTMStateTuple(x2, x2)
        cell = AttentionWrapper(lstm, attention_mechanism, output_attention=False)
        cell_state = cell.zero_state(dtype=tf.float32, batch_size=test_batch_size * beam_width)
        cell_state = cell_state.clone(cell_state=encoder_state2, attention=tiled_first_attention)
        infer_decoder = BeamSearchDecoder(cell, embedding=label_embeddings, start_tokens=[GO] * test_len, end_token=EOS,
                                          initial_state=cell_state, beam_width=beam_width, output_layer=output_l)
        decoder_outputs_infer, decoder_state_infer, decoder_seq_infer = dynamic_decode(infer_decoder, maximum_iterations=4)
        return decoder_outputs_train, decoder_outputs_infer, decoder_state_infer 

def state_size(self):
        return (LSTMStateTuple(self._num_units, self._num_units)
                if self._state_is_tuple else 2 * self._num_units) 

def state_size(self):
        return (LSTMStateTuple(self._num_units, self._num_units)
                if self._state_is_tuple else 2 * self._num_units) 

def _get_state_names(cell):
  """Gets the state names for an `RNNCell`.

  Args:
    cell: A `RNNCell` to be used in the RNN.

  Returns:
    State names in the form of a string, a list of strings, or a list of
    string pairs, depending on the type of `cell.state_size`.

  Raises:
    TypeError: If cell.state_size is of type TensorShape.
  """
  state_size = cell.state_size
  if isinstance(state_size, tensor_shape.TensorShape):
    raise TypeError('cell.state_size of type TensorShape is not supported.')
  if isinstance(state_size, int):
    return '{}_{}'.format(rnn_common.RNNKeys.STATE_PREFIX, 0)
  if isinstance(state_size, rnn_cell.LSTMStateTuple):
    return [
        '{}_{}_c'.format(rnn_common.RNNKeys.STATE_PREFIX, 0),
        '{}_{}_h'.format(rnn_common.RNNKeys.STATE_PREFIX, 0),
    ]
  if isinstance(state_size[0], rnn_cell.LSTMStateTuple):
    return [[
        '{}_{}_c'.format(rnn_common.RNNKeys.STATE_PREFIX, i),
        '{}_{}_h'.format(rnn_common.RNNKeys.STATE_PREFIX, i),
    ] for i in range(len(state_size))]
  return [
      '{}_{}'.format(rnn_common.RNNKeys.STATE_PREFIX, i)
      for i in range(len(state_size))] 

def batch_gather(batch_dict, indices, _top=True):
    """
    Gathers experiences from processed batch according to specified indices.

    Args:
        batch_dict:     batched data dictionary
        indices:        array-like, indices to gather
        _top:           internal

    Returns:
        batched data of same structure as dict

    """
    batch = {}

    if isinstance(batch_dict, dict):
        for key, value in batch_dict.items():
            batch[key] = batch_gather(value, indices, False)

    elif isinstance(batch_dict, LSTMStateTuple):
        c = batch_gather(batch_dict[0], indices, False)
        h = batch_gather(batch_dict[1], indices, False)
        batch = LSTMStateTuple(c=c, h=h)

    elif isinstance(batch_dict, tuple):
        batch = tuple([batch_gather(struct, indices, False) for struct in batch_dict])

    else:
        batch = np.take(batch_dict, indices=indices, axis=0, mode='wrap')

    if _top:
        # Mind shape inference:
        batch['batch_size'] = indices.shape[0]

    return batch 

def pop_frame(self, idx, _struct=None):
        """
        Pops single experience from rollout.

        Args:
            idx:    experience position

        Returns:
            frame as [nested] dictionary
        """
        # No idx range checks here!
        if _struct is None:
            _struct = self

        if isinstance(_struct, dict) or type(_struct) == type(self):
            frame = {}
            for key, value in _struct.items():
                frame[key] = self.pop_frame(idx, value)
            return frame

        elif isinstance(_struct, tuple):
            return tuple([self.pop_frame(idx, value) for value in _struct])

        elif isinstance(_struct, LSTMStateTuple):
            return LSTMStateTuple(self.pop_frame(idx, _struct[0]), self.pop_frame(idx, _struct[1]))

        else:
            return _struct.pop(idx) 

def state_size(self):
        return rnn.LSTMStateTuple(self._num_units, self._num_units) 

def __call__(self, inputs, state, scope=None):
        with tf.variable_scope(scope or type(self).__name__, reuse=self._reuse):
            c, h = state
            input_size = inputs.get_shape().as_list()[1]
            W_xh = tf.get_variable('W_xh',
                                   [input_size, 4 * self._num_units],
                                   initializer=orthogonal_initializer())
            W_hh = tf.get_variable('W_hh',
                                   [self._num_units, 4 * self._num_units],
                                   initializer=bn_lstm_identity_initializer(0.95))
            bias = tf.get_variable('bias', [4 * self._num_units])

            xh = tf.matmul(inputs, W_xh)
            hh = tf.matmul(h, W_hh)

            bn_xh = batch_norm(xh, self._is_training)
            bn_hh = batch_norm(hh, self._is_training)

            hidden = bn_xh + bn_hh + bias

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

            new_c = (c * sigmoid(f + self._forget_bias) + sigmoid(i) * self._activation(j))
            bn_new_c = batch_norm(new_c, 'c', self._is_training)
            new_h = self._activation(bn_new_c) * sigmoid(o)
            new_state = rnn.LSTMStateTuple(new_c, new_h)

            return new_h, new_state 

def state_size(self):
        return rnn.LSTMStateTuple(self._num_units, self._num_units) 

def __call__(self, inputs, state, scope=None):
        with tf.variable_scope(scope or type(self).__name__, reuse=self._reuse):
            c, h = state
            input_size = inputs.get_shape().as_list()[1]
            W_xh = tf.get_variable('W_xh',
                                   [input_size, 4 * self._num_units],
                                   initializer=orthogonal_initializer())
            W_hh = tf.get_variable('W_hh',
                                   [self._num_units, 4 * self._num_units],
                                   initializer=bn_lstm_identity_initializer(0.95))
            bias = tf.get_variable('bias', [4 * self._num_units])

            xh = tf.matmul(inputs, W_xh)
            hh = tf.matmul(h, W_hh)

            bn_xh = batch_norm(xh, self._is_training)
            bn_hh = batch_norm(hh, self._is_training)

            hidden = bn_xh + bn_hh + bias

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

            new_c = (c * sigmoid(f + self._forget_bias) + sigmoid(i) * self._activation(j))
            bn_new_c = batch_norm(new_c, 'c', self._is_training)
            new_h = self._activation(bn_new_c) * sigmoid(o)
            new_state = rnn.LSTMStateTuple(new_c, new_h)

            return new_h, new_state 

def state_size(self):
        return rnn.LSTMStateTuple(self._num_units, self._num_units) 

def __call__(self, inputs, state, scope=None):
        with tf.variable_scope(scope or type(self).__name__, reuse=self._reuse):
            c, h = state
            input_size = inputs.get_shape().as_list()[1]
            W_xh = tf.get_variable('W_xh',
                                   [input_size, 4 * self._num_units],
                                   initializer=orthogonal_initializer())
            W_hh = tf.get_variable('W_hh',
                                   [self._num_units, 4 * self._num_units],
                                   initializer=bn_lstm_identity_initializer(0.95))
            bias = tf.get_variable('bias', [4 * self._num_units])

            xh = tf.matmul(inputs, W_xh)
            hh = tf.matmul(h, W_hh)

            bn_xh = batch_norm(xh, self._is_training)
            bn_hh = batch_norm(hh, self._is_training)

            hidden = bn_xh + bn_hh + bias

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

            new_c = (c * sigmoid(f + self._forget_bias) + sigmoid(i) * self._activation(j))
            bn_new_c = batch_norm(new_c, 'c', self._is_training)
            new_h = self._activation(bn_new_c) * sigmoid(o)
            new_state = rnn.LSTMStateTuple(new_c, new_h)

            return new_h, new_state 

def state_size(self):
        return (LSTMStateTuple(self._num_units, self._num_units)
                if self._state_is_tuple else 2 * self._num_units) 

def __call__(self, inputs, state, scope=None):
        """Long short-term memory cell (LSTM)."""
        with tf.variable_scope(scope or type(self).__name__):  # "BasicLSTMCell"
            # Parameters of gates are concatenated into one multiply for efficiency.
            if self._state_is_tuple:
                c, h = state
            else:
                c, h = tf.split(axis=1, num_or_size_splits=2, value=state)

            batch_size = tf.shape(inputs)[0]
            inputs = tf.reshape(inputs, [batch_size, self.height, self.width, 1])
            c = tf.reshape(c, [batch_size, self.height, self.width, self.num_features])
            h = tf.reshape(h, [batch_size, self.height, self.width, self.num_features])

            concat = _conv_linear([inputs, h], self.filter_size, self.num_features * 4, True)

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

            new_c = (c * tf.nn.sigmoid(f + self._forget_bias) + tf.nn.sigmoid(i) *
                     self._activation(j))
            new_h = self._activation(new_c) * tf.nn.sigmoid(o)

            new_h = tf.reshape(new_h, [batch_size, self._num_units]) 
            new_c = tf.reshape(new_c, [batch_size, self._num_units])

            if self._state_is_tuple:
                new_state = LSTMStateTuple(new_c, new_h)
            else:
                new_state = tf.concat(axis=1, values=[new_c, new_h])
            
            return new_h, new_state 

def single_layer_shape(self):
        if 'LSTM' in self._base_cell:
            return LSTMStateTuple(c=tf.TensorShape([None, self._state_size]),
                                  h=tf.TensorShape([None, self._state_size]))
        else:
            return tf.TensorShape([None, self._state_size]) 

def wrapper(self, state):
        """Some RNN states are wrapped in namedtuples. 
        (TensorFlow decision, definitely not mine...). 
        
        This is here for derived classes to specify their wrapper state. 
        Some examples: LSTMStateTuple and AttentionWrapperState.
        
        Args:
            state: tensor state tuple, will be unpacked into the wrapper tuple.
        """
        if self._wrapper is None:
            return state
        else:
            return self._wrapper(*state) 

def __init__(self, cell, output_size):
        """Create a cell with with a linear encoder in space.

        Args:
          cell: an RNNCell. The input is passed through a linear layer.

        Raises:
          TypeError: if cell is not an RNNCell.
        """
        if not isinstance(cell, RNNCell):
            raise TypeError("The parameter cell is not a RNNCell.")

        self._cell = cell

        print('output_size = {0}'.format(output_size))
        print(' state_size = {0}'.format(self._cell.state_size))

        # Tuple if multi-rnn
        if isinstance(self._cell.state_size, tuple):

            # Fine if GRU...
            insize = self._cell.state_size[-1]

            # LSTMStateTuple if LSTM
            if isinstance(insize, LSTMStateTuple):
                insize = insize.h

        else:
            # Fine if not multi-rnn
            insize = self._cell.state_size

        self.w_out = tf.get_variable("proj_w_out",
                                     [insize, output_size],
                                     dtype=tf.float32,
                                     initializer=tf.random_uniform_initializer(minval=-0.04, maxval=0.04))
        self.b_out = tf.get_variable("proj_b_out", [output_size],
                                     dtype=tf.float32,
                                     initializer=tf.random_uniform_initializer(minval=-0.04, maxval=0.04))

        self.linear_output_size = output_size 

def __call__(self, states):
        state_transfered = self.init_transfer(states)
        if self.drnn_state_is_tuple:
            initial_state = LSTMStateTuple(h=state_transfered[:,:self.drnn_state_size],
                                           c=state_transfered[:,self.drnn_state_size:])
            return tuple(initial_state for _ in range(self.drnn_num_layers))

        return tuple(state_transfered for _ in range(self.num_layers)) 

def _beam_step(self, time, inputs, state):
        cell_output_new, logits, attention_scores, attention_context, cell_states = \
            self.compute_output(inputs, state)

        sample_ids, beam_ids, log_probs = \
            self.feedback.sample(logits=logits,
                                 log_probs=state.log_probs,
                                 prev_finished=state.finished,
                                 time=time)

        # attention_scores = tf.gather(attention_scores, beam_ids)
        attention_context = tf.gather(attention_context, beam_ids)

        def gather_state(state, idx):
            if isinstance(state, tf.Tensor):
                return tf.gather(state, idx)
            elif isinstance(state, tf.contrib.rnn.LSTMStateTuple):
                c = tf.gather(state.c, idx)
                h = tf.gather(state.h, idx)
                return tf.contrib.rnn.LSTMStateTuple(c, h)
            else:
                raise ValueError(
                    "Unrecognized state type: %s" % str(type(state)))

        next_cell_states = tuple(
            [gather_state(state, beam_ids) for state in cell_states])

        outputs = self.output_tuple(
            logits=logits,
            predicted_ids=sample_ids,
            beam_ids=beam_ids)

        finished, next_inputs = self.feedback.next_inputs(time=time,
                                                          sample_ids=sample_ids)

        next_inputs = tf.concat([next_inputs, attention_context], 1)
        next_state = self.state_tuple(cell_states=next_cell_states,
                                      log_probs=log_probs,
                                      finished=finished)

        return (outputs, next_state, next_inputs, finished) 

def _make_encoder(self):
        """Create the encoder"""
        inputs = layers.embed_sequence(
            self.X,
            vocab_size=self.vocab_size,
            embed_dim=self.embed_dim,
            scope='embed')

        # Project to correct dimensions
        # b/c the bidirectional RNN's forward and backward
        # outputs are concatenated, the size will be 2x,
        # so halve the hidden sizes here to compensate
        inputs = tf.layers.dense(inputs, self.hidden_size//2)

        cell_fw = rnn.MultiRNNCell([
            self._make_cell(self.hidden_size//2) for _ in range(self.depth)
        ])
        cell_bw = rnn.MultiRNNCell([
            self._make_cell(self.hidden_size//2) for _ in range(self.depth)
        ])
        encoder_outputs, encoder_final_state = tf.nn.bidirectional_dynamic_rnn(
            cell_fw=cell_fw, cell_bw=cell_bw, sequence_length=self.sequence_length,
            inputs=inputs, dtype=tf.float32)

        # Concat forward and backward outputs
        encoder_outputs = tf.concat(encoder_outputs, 2)

        # Concat forward and backward layer states
        encoder_fw_states, encoder_bw_states = encoder_final_state
        encoder_final_state = []
        for fw, bw in zip(encoder_fw_states, encoder_bw_states):
            c = tf.concat([fw.c, bw.c], 1)
            h = tf.concat([fw.h, bw.h], 1)
            encoder_final_state.append(rnn.LSTMStateTuple(c=c, h=h))
        return encoder_outputs, encoder_final_state 

def _project_lstm_state_tuple(state_tuple, num_units):

    state_proj_layer = Dense(num_units, name='state_projection', use_bias=False)

    cat_c = tf.concat([state.c for state in state_tuple], axis=-1)
    cat_h = tf.concat([state.h for state in state_tuple], axis=-1)

    proj_c = state_proj_layer(cat_c)
    proj_h = state_proj_layer(cat_h)

    projected_state = tf.contrib.rnn.LSTMStateTuple(c=proj_c, h=proj_h)

    return projected_state