Python tensorflow.contrib.rnn.LSTMStateTuple() Examples

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 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 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 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 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 __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 __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 initial_states_tuple(self):
        """
        Create the initial state tensors for the individual RNN cells.
        
        If no initial state vector was passed to this RNN, all initial states are set to be zero. Otherwise, the initial
        state vector is split into a possibly nested tuple of tensors according to the RNN architecture. The return
        value of this function is structured in such a way that it can be passed to the `initial_state` parameter of the
        RNN functions in `tf.contrib.rnn`.
        
        Returns
        -------
        tuple of tf.Tensor
            A possibly nested tuple of initial state tensors for the RNN cells
        """
        if self.initial_state is None:
            initial_states = tf.zeros(shape=[self.batch_size, self.state_size], dtype=tf.float32)
        else:
            initial_states = self.initial_state

        initial_states = tuple(tf.split(initial_states, self.num_layers, axis=1))

        if self.bidirectional:
            initial_states = tuple([tf.split(x, 2, axis=1) for x in initial_states])
            initial_states_fw, initial_states_bw = zip(*initial_states)

            if self.cell_type == CellType.LSTM:
                initial_states_fw = tuple([LSTMStateTuple(*tf.split(lstm_state, 2, axis=1))
                                           for lstm_state in initial_states_fw])
                initial_states_bw = tuple([LSTMStateTuple(*tf.split(lstm_state, 2, axis=1))
                                           for lstm_state in initial_states_bw])

            initial_states = (initial_states_fw, initial_states_bw)
        else:
            if self.cell_type == CellType.LSTM:
                initial_states = tuple([LSTMStateTuple(*tf.split(lstm_state, 2, axis=1))
                                        for lstm_state in initial_states])

        return initial_states 

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

def __init__(self, ob_space, ac_space, lstm_size=256, use_categorical_max=False, **kwargs):
        self.x = x = tf.placeholder(tf.float32, [None] + list(ob_space))

        rank = len(ob_space)

        if rank == 3: # pixel input
            for i in range(4):
                x = tf.nn.elu(conv2d(x, 32, "l{}".format(i + 1), [3, 3], [2, 2]))
        elif rank == 1: # plain features
            #x = tf.nn.elu(linear(x, 256, "l1", normalized_columns_initializer(0.01)))
            pass
        else:
            raise TypeError("observation space must have rank 1 or 3, got %d" % rank)

        # introduce a "fake" batch dimension of 1 after flatten so that we can do LSTM over time dim
        x = tf.expand_dims(flatten(x), [0])

        size = lstm_size
        lstm = rnn.BasicLSTMCell(size, state_is_tuple=True)
        self.state_size = lstm.state_size
        step_size = tf.shape(self.x)[:1]

        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, [1, lstm.state_size.c])
        h_in = tf.placeholder(tf.float32, [1, lstm.state_size.h])
        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_c, lstm_h = lstm_state
        x = tf.reshape(lstm_outputs, [-1, size])
        self.logits = linear(x, ac_space, "action", normalized_columns_initializer(0.01))
        self.vf = tf.reshape(linear(x, 1, "value", normalized_columns_initializer(1.0)), [-1])
        self.state_out = [lstm_c[:1, :], lstm_h[:1, :]]
        self.sample = categorical_max(self.logits, ac_space)[0, :] \
            if use_categorical_max else categorical_sample(self.logits, ac_space)[0, :]
        self.var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, tf.get_variable_scope().name) 

def pass_messages(self):
        with tf.name_scope('pass_messages') as scope:
            denom = tf.sqrt(tf.cast(self.opts.d, tf.float32))

            L_output = tf.tile(tf.div(self.L_init, denom), [self.n_lits, 1])
            C_output = tf.tile(tf.div(self.C_init, denom), [self.n_clauses, 1])

            L_state = LSTMStateTuple(h=L_output, c=tf.zeros([self.n_lits, self.opts.d]))
            C_state = LSTMStateTuple(h=C_output, c=tf.zeros([self.n_clauses, self.opts.d]))

            _, L_state, C_state = tf.while_loop(self.while_cond, self.while_body, [0, L_state, C_state])

        self.final_lits = L_state.h
        self.final_clauses = C_state.h 

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) 

def convert_to_state(self, variables):
        if len(variables) != 2:
            raise ValueError()
        return LSTMStateTuple(variables[0], variables[1]) 

def convert_to_state(self, variables):
        if len(variables) != 2:
            raise ValueError()
        return LSTMStateTuple(variables[0], variables[1]) 

def convert_to_state(self, variables):
        if len(variables) != 2:
            raise ValueError()
        return LSTMStateTuple(variables[0], variables[1]) 

def state_size_flat(self):
    return contrib_rnn.LSTMStateTuple([self._param_count], [self._param_count]) 

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 batch is None:

            init_states = self.session.run(self.initial_state, {self.batch_size: batchsize})
            predicted_outputs = self.session.run(self.predicted_outputs,
                                                 {self.rnn_outputs: init_states.h[None, :, :]})
            return predicted_outputs, init_states

        else:
            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 __init__(self, ob_space, ac_space, lstm_size=256, **kwargs):
        self.x = x = tf.placeholder(tf.float32, [None] + list(ob_space))

        rank = len(ob_space)

        if rank == 3: # pixel input
            for i in range(4):
                x = tf.nn.elu(conv2d(x, 32, "l{}".format(i + 1), [3, 3], [2, 2]))
        elif rank == 1: # plain features
            #x = tf.nn.elu(linear(x, 256, "l1", normalized_columns_initializer(0.01)))
            pass
        else:
            raise TypeError("observation space must have rank 1 or 3, got %d" % rank)

        # introduce a "fake" batch dimension of 1 after flatten so that we can do LSTM over time dim
        x = tf.expand_dims(flatten(x), [0])

        size = lstm_size
        lnlstm = rnn.LayerNormBasicLSTMCell(size)
        self.state_size = lnlstm.state_size
        step_size = tf.shape(self.x)[:1]

        c_init = np.zeros((1, lnlstm.state_size.c), np.float32)
        h_init = np.zeros((1, lnlstm.state_size.h), np.float32)
        self.state_init = [c_init, h_init]
        c_in = tf.placeholder(tf.float32, [1, lnlstm.state_size.c])
        h_in = tf.placeholder(tf.float32, [1, lnlstm.state_size.h])
        self.state_in = [c_in, h_in]

        state_in = rnn.LSTMStateTuple(c_in, h_in)
        lstm_outputs, lstm_state = tf.nn.dynamic_rnn(
            lnlstm, x, initial_state=state_in, sequence_length=step_size,
            time_major=False)
        lstm_c, lstm_h = lstm_state
        x = tf.reshape(lstm_outputs, [-1, size])
        self.logits = linear(x, ac_space, "action", normalized_columns_initializer(0.01))
        self.vf = tf.reshape(linear(x, 1, "value", normalized_columns_initializer(1.0)), [-1])
        self.state_out = [lstm_c[:1, :], lstm_h[:1, :]]
        self.sample = categorical_sample(self.logits, ac_space)[0, :]
        self.var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, tf.get_variable_scope().name) 

def __init__(self, ob_space, ac_space, lstm_size=256, use_categorical_max=False, **kwargs):
        self.x = x = tf.placeholder(tf.float32, [None] + list(ob_space))

        rank = len(ob_space)

        if rank == 3: # pixel input
            for i in range(4):
                x = tf.nn.elu(conv2d(x, 32, "l{}".format(i + 1), [3, 3], [2, 2]))
        elif rank == 1: # plain features
            #x = tf.nn.elu(linear(x, 256, "l1", normalized_columns_initializer(0.01)))
            pass
        else:
            raise TypeError("observation space must have rank 1 or 3, got %d" % rank)

        # introduce a "fake" batch dimension of 1 after flatten so that we can do LSTM over time dim
        x = tf.expand_dims(flatten(x), [0])

        size = lstm_size
        lstm = rnn.BasicLSTMCell(size, state_is_tuple=True)
        self.state_size = lstm.state_size
        step_size = tf.shape(self.x)[:1]

        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, [1, lstm.state_size.c])
        h_in = tf.placeholder(tf.float32, [1, lstm.state_size.h])
        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_c, lstm_h = lstm_state
        x = tf.reshape(lstm_outputs, [-1, size])
        self.logits = linear(x, ac_space, "action", normalized_columns_initializer(0.01))
        self.vf = tf.reshape(linear(x, 1, "value", normalized_columns_initializer(1.0)), [-1])
        self.state_out = [lstm_c[:1, :], lstm_h[:1, :]]
        self.sample = categorical_max(self.logits, ac_space)[0, :] \
            if use_categorical_max else categorical_sample(self.logits, ac_space)[0, :]
        self.var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, tf.get_variable_scope().name) 

def __init__(self, ob_space, ac_space, lstm_size=256, **kwargs):
        self.x = x = tf.placeholder(tf.float32, [None] + list(ob_space))

        rank = len(ob_space)

        if rank == 3: # pixel input
            for i in range(4):
                x = tf.nn.elu(conv2d(x, 32, "l{}".format(i + 1), [3, 3], [2, 2]))
        elif rank == 1: # plain features
            #x = tf.nn.elu(linear(x, 256, "l1", normalized_columns_initializer(0.01)))
            pass
        else:
            raise TypeError("observation space must have rank 1 or 3, got %d" % rank)

        # introduce a "fake" batch dimension of 1 after flatten so that we can do LSTM over time dim
        x = tf.expand_dims(flatten(x), [0])

        size = lstm_size
        lnlstm = rnn.LayerNormBasicLSTMCell(size)
        self.state_size = lnlstm.state_size
        step_size = tf.shape(self.x)[:1]

        c_init = np.zeros((1, lnlstm.state_size.c), np.float32)
        h_init = np.zeros((1, lnlstm.state_size.h), np.float32)
        self.state_init = [c_init, h_init]
        c_in = tf.placeholder(tf.float32, [1, lnlstm.state_size.c])
        h_in = tf.placeholder(tf.float32, [1, lnlstm.state_size.h])
        self.state_in = [c_in, h_in]

        state_in = rnn.LSTMStateTuple(c_in, h_in)
        lstm_outputs, lstm_state = tf.nn.dynamic_rnn(
            lnlstm, x, initial_state=state_in, sequence_length=step_size,
            time_major=False)
        lstm_c, lstm_h = lstm_state
        x = tf.reshape(lstm_outputs, [-1, size])
        self.logits = linear(x, ac_space, "action", normalized_columns_initializer(0.01))
        self.vf = tf.reshape(linear(x, 1, "value", normalized_columns_initializer(1.0)), [-1])
        self.state_out = [lstm_c[:1, :], lstm_h[:1, :]]
        self.sample = categorical_sample(self.logits, ac_space)[0, :]
        self.var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, tf.get_variable_scope().name) 

def state_size(self):
    return contrib_rnn.LSTMStateTuple(self._output_size + [self._num_units],
                                      self._output_size + [self._num_units]) 

def cudnn_lstm_state_to_state_tuples(cudnn_lstm_state):
  """Convert CudnnLSTM format to tuple of LSTMStateTuples."""
  h, c = cudnn_lstm_state
  return tuple(
      rnn.LSTMStateTuple(h=h_i, c=c_i)
      for h_i, c_i in zip(tf.unstack(h), tf.unstack(c))) 

def __call__(self, inputs, state, scope=None):
        """Long short-term memory cell (LSTM).
        @param: inputs (batch,n)
        @param state: the states and hidden unit of the two cells
        """
        with tf.variable_scope(scope or type(self).__name__):
            c1, c2, h1, h2 = state

            # change bias argument to False since LN will add bias via shift
            concat = _linear([inputs, h1, h2], 5 * self._num_units, False)

            i, j, f1, f2, o = tf.split(value=concat, num_or_size_splits=5, axis=1)

            # add layer normalization to each gate
            i = ln(i, scope='i/')
            j = ln(j, scope='j/')
            f1 = ln(f1, scope='f1/')
            f2 = ln(f2, scope='f2/')
            o = ln(o, scope='o/')

            new_c = (c1 * tf.nn.sigmoid(f1 + self._forget_bias) +
                     c2 * tf.nn.sigmoid(f2 + self._forget_bias) + tf.nn.sigmoid(i) *
                     self._activation(j))

            # add layer_normalization in calculation of new hidden state
            new_h = self._activation(ln(new_c, scope='new_h/')) * tf.nn.sigmoid(o)
            new_state = LSTMStateTuple(new_c, new_h)

            return new_h, new_state 

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

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 __init__(self, ob_space, ac_space):
        self.x = x = tf.placeholder(tf.float32, [None] + list(ob_space))

        for i in range(4):
            x = tf.nn.elu(conv2d(x, 32, "l{}".format(i + 1), [3, 3], [2, 2]))
        # introduce a "fake" batch dimension of 1 after flatten so that we can do LSTM over time dim
        x = tf.expand_dims(flatten(x), [0])

        size = 256
        if use_tf100_api:
            lstm = rnn.BasicLSTMCell(size, state_is_tuple=True)
        else:
            lstm = rnn.rnn_cell.BasicLSTMCell(size, state_is_tuple=True)
        self.state_size = lstm.state_size
        step_size = tf.shape(self.x)[:1]

        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, [1, lstm.state_size.c])
        h_in = tf.placeholder(tf.float32, [1, lstm.state_size.h])
        self.state_in = [c_in, h_in]

        if use_tf100_api:
            state_in = rnn.LSTMStateTuple(c_in, h_in)
        else:
            state_in = rnn.rnn_cell.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_c, lstm_h = lstm_state
        x = tf.reshape(lstm_outputs, [-1, size])
        self.logits = linear(x, ac_space, "action", normalized_columns_initializer(0.01))
        self.vf = tf.reshape(linear(x, 1, "value", normalized_columns_initializer(1.0)), [-1])
        self.state_out = [lstm_c[:1, :], lstm_h[:1, :]]
        self.sample = categorical_sample(self.logits, ac_space)[0, :]
        self.var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, tf.get_variable_scope().name) 

def __call__(self, inputs, state, scope=None):
        with tf.name_scope('context_cell'):
            cell_output, cell_state = self._cell(
                inputs, state.cell_state, scope)
            context_len = tf.shape(self._context)[1]
            context_input = self._context[:, state.time % context_len, :]

            # First we apply context weighting to the cell output
            output_size = self._cell.output_size
            context_output = tf.concat([cell_output, context_input], axis=1)
            context_layer = tf.layers.dense(
                context_output, output_size, name='output_ctx')
            context_output = self._activation(context_layer)

            # Next we add context weighting to the cell state
            if isinstance(self._cell.state_size, LSTMStateTuple):
                # For LSTM cells we only apply context to the cell state
                state_size = self._cell.state_size.c
                context_state = tf.concat(
                    [cell_state.c, context_input], axis=1)
                context_layer = tf.layers.dense(
                    context_state, state_size, name='state_ctx')
                context_state = LSTMStateTuple(
                    c=self._activation(context_layer),
                    h=context_output
                )
            else:
                state_size = self._cell.state_size
                context_state = tf.concat([cell_state, context_input], axis=1)
                context_layer = tf.layers.dense(
                    context_state, state_size, name='state_ctx')
                context_state = self._activation(context_layer)

            return context_output, ContextState(
                time=state.time + 1,
                cell_state=context_state
            )