# Path: seya/utils.py
# def rnn_states(step_function, inputs, initial_states,
#                go_backwards=False, masking=True):
#     '''Iterates over the time dimension of a tensor.
#     Parameters
#     ----------
#     inputs: tensor of temporal data of shape (samples, time, ...)
#         (at least 3D).
#     step_function:
#         Parameters:
#             input: tensor with shape (samples, ...) (no time dimension),
#                 representing input for the batch of samples at a certain
#                 time step.
#             states: list of tensors.
#         Returns:
#             output: tensor with shape (samples, ...) (no time dimension),
#             new_states: list of tensors, same length and shapes
#                 as 'states'.
#     initial_states: tensor with shape (samples, ...) (no time dimension),
#         containing the initial values for the states used in
#         the step function.
#     go_backwards: boolean. If True, do the iteration over
#         the time dimension in reverse order.
#     masking: boolean. If true, any input timestep inputs[s, i]
#         that is all-zeros will be skipped (states will be passed to
#         the next step unchanged) and the corresponding output will
#         be all zeros.
#     Returns
#     -------
#     A tuple (last_output, outputs, new_states).
#         last_output: the latest output of the rnn, of shape (samples, ...)
#         outputs: tensor with shape (samples, time, ...) where each
#             entry outputs[s, t] is the output of the step function
#             at time t for sample s.
#         new_states: list of tensors, latest states returned by
#             the step function, of shape (samples, ...).
#     '''
#     inputs = inputs.dimshuffle((1, 0, 2))
# 
#     def _step(input, *states):
#         output, new_states = step_function(input, states)
#         if masking:
#             # if all-zero input timestep, return
#             # all-zero output and unchanged states
#             switch = T.any(input, axis=-1, keepdims=True)
#             output = T.switch(switch, output, 0. * output)
#             return_states = []
#             for state, new_state in zip(states, new_states):
#                 return_states.append(T.switch(switch, new_state, state))
#             return [output] + return_states
#         else:
#             return [output] + new_states
# 
#     results, _ = theano.scan(
#         _step,
#         sequences=inputs,
#         outputs_info=[None] + initial_states,
#         go_backwards=go_backwards)
# 
#     # deal with Theano API inconsistency
#     if type(results) is list:
#         states = results[1:]
#     else:
#         states = []
# 
#     return states


import numpy as np
import theano
import theano.tensor as T
from keras.layers.recurrent import Recurrent, GRU, LSTM
from keras import backend as K
from seya.utils import rnn_states
        [memory, read_address, write_address, rnn_state] = F(x)

        if inner_rnn == "lstm" use it as
        [memory, read_address, write_address, rnn_cell, rnn_state] = F(x)

        """
        # input shape: (nb_samples, time (padded with zeros), input_dim)
        X = self.get_input(train)
        assert K.ndim(X) == 3
        if K._BACKEND == 'tensorflow':
            if not self.input_shape[1]:
                raise Exception('When using TensorFlow, you should define ' +
                                'explicitely the number of timesteps of ' +
                                'your sequences. Make sure the first layer ' +
                                'has a "batch_input_shape" argument ' +
                                'including the samples axis.')

        mask = self.get_output_mask(train)
        if mask:
            # apply mask
            X *= K.cast(K.expand_dims(mask), X.dtype)
            masking = True
        else:
            masking = False

        if self.stateful:
            initial_states = self.states
        else:
            initial_states = self.get_initial_states(X)


# Path: seya/utils.py
# def rnn_states(step_function, inputs, initial_states,
#                go_backwards=False, masking=True):
#     '''Iterates over the time dimension of a tensor.
#     Parameters
#     ----------
#     inputs: tensor of temporal data of shape (samples, time, ...)
#         (at least 3D).
#     step_function:
#         Parameters:
#             input: tensor with shape (samples, ...) (no time dimension),
#                 representing input for the batch of samples at a certain
#                 time step.
#             states: list of tensors.
#         Returns:
#             output: tensor with shape (samples, ...) (no time dimension),
#             new_states: list of tensors, same length and shapes
#                 as 'states'.
#     initial_states: tensor with shape (samples, ...) (no time dimension),
#         containing the initial values for the states used in
#         the step function.
#     go_backwards: boolean. If True, do the iteration over
#         the time dimension in reverse order.
#     masking: boolean. If true, any input timestep inputs[s, i]
#         that is all-zeros will be skipped (states will be passed to
#         the next step unchanged) and the corresponding output will
#         be all zeros.
#     Returns
#     -------
#     A tuple (last_output, outputs, new_states).
#         last_output: the latest output of the rnn, of shape (samples, ...)
#         outputs: tensor with shape (samples, time, ...) where each
#             entry outputs[s, t] is the output of the step function
#             at time t for sample s.
#         new_states: list of tensors, latest states returned by
#             the step function, of shape (samples, ...).
#     '''
#     inputs = inputs.dimshuffle((1, 0, 2))
# 
#     def _step(input, *states):
#         output, new_states = step_function(input, states)
#         if masking:
#             # if all-zero input timestep, return
#             # all-zero output and unchanged states
#             switch = T.any(input, axis=-1, keepdims=True)
#             output = T.switch(switch, output, 0. * output)
#             return_states = []
#             for state, new_state in zip(states, new_states):
#                 return_states.append(T.switch(switch, new_state, state))
#             return [output] + return_states
#         else:
#             return [output] + new_states
# 
#     results, _ = theano.scan(
#         _step,
#         sequences=inputs,
#         outputs_info=[None] + initial_states,
#         go_backwards=go_backwards)
# 
#     # deal with Theano API inconsistency
#     if type(results) is list:
#         states = results[1:]
#     else:
#         states = []
# 
#     return states


# Path: seya/layers/ntm.py
import numpy as np
import theano
import theano.tensor as T
from keras.layers.recurrent import Recurrent, GRU, LSTM
from keras import backend as K
from seya.utils import rnn_states

        [memory, read_address, write_address, rnn_state] = F(x)

        if inner_rnn == "lstm" use it as
        [memory, read_address, write_address, rnn_cell, rnn_state] = F(x)

        """
        # input shape: (nb_samples, time (padded with zeros), input_dim)
        X = self.get_input(train)
        assert K.ndim(X) == 3
        if K._BACKEND == 'tensorflow':
            if not self.input_shape[1]:
                raise Exception('When using TensorFlow, you should define ' +
                                'explicitely the number of timesteps of ' +
                                'your sequences. Make sure the first layer ' +
                                'has a "batch_input_shape" argument ' +
                                'including the samples axis.')

        mask = self.get_output_mask(train)
        if mask:
            # apply mask
            X *= K.cast(K.expand_dims(mask), X.dtype)
            masking = True
        else:
            masking = False

        if self.stateful:
            initial_states = self.states
        else:
            initial_states = self.get_initial_states(X)


        states = rnn_states(self.step, X, initial_states,