rnn_cell.py

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import numpy as np
import random

from caffe2.python.attention import (
    AttentionType,
    apply_regular_attention,
    apply_recurrent_attention,
)
from caffe2.python import core, recurrent, workspace
from caffe2.python.cnn import CNNModelHelper


class RNNCell(object):
    '''
    Base class for writing recurrent / stateful operations.

    One needs to implement 3 methods: _apply, prepare_input and get_state_names.
    As a result base class will provice apply_over_sequence method, which
    allows you to apply recurrent operations over a sequence of any length.
    '''
    def __init__(self, name):
        self.name = name
        self.recompute_blobs = []

    def scope(self, name):
        return self.name + '/' + name if self.name is not None else name

    def apply_over_sequence(
        self,
        model,
        inputs,
        seq_lengths,
        initial_states,
        outputs_with_grads=None,
    ):
        preprocessed_inputs = self.prepare_input(model, inputs)
        step_model = CNNModelHelper(name=self.name, param_model=model)
        input_t, timestep = step_model.net.AddScopedExternalInputs(
            'input_t',
            'timestep',
        )
        states_prev = step_model.net.AddScopedExternalInputs(*[
            s + '_prev' for s in self.get_state_names()
        ])
        states = self._apply(
            model=step_model,
            input_t=input_t,
            seq_lengths=seq_lengths,
            states=states_prev,
            timestep=timestep,
        )
        return recurrent.recurrent_net(
            net=model.net,
            cell_net=step_model.net,
            inputs=[(input_t, preprocessed_inputs)],
            initial_cell_inputs=zip(states_prev, initial_states),
            links=dict(zip(states_prev, states)),
            timestep=timestep,
            scope=self.name,
            outputs_with_grads=(
                outputs_with_grads
                if outputs_with_grads is not None
                else self.get_outputs_with_grads()
            ),
            recompute_blobs_on_backward=self.recompute_blobs,
        )

    def apply(self, model, input_t, seq_lengths, states, timestep):
        input_t = self.prepare_input(model, input_t)
        return self._apply(model, input_t, seq_lengths, states, timestep)

    def _apply(self, model, input_t, seq_lengths, states, timestep):
        '''
        A single step of a recurrent network.

        model: CNNModelHelper object new operators would be added to

        input_blob: single input with shape (1, batch_size, input_dim)

        seq_lengths: blob containing sequence lengths which would be passed to
        LSTMUnit operator

        states: previous recurrent states

        timestep: current recurrent iteration. Could be used together with
        seq_lengths in order to determine, if some shorter sequences
        in the batch have already ended.
        '''
        raise NotImplementedError('Abstract method')

    def prepare_input(self, model, input_blob):
        '''
        If some operations in _apply method depend only on the input,
        not on recurrent states, they could be computed in advance.

        model: CNNModelHelper object new operators would be added to

        input_blob: either the whole input sequence with shape
        (sequence_length, batch_size, input_dim) or a single input with shape
        (1, batch_size, input_dim).
        '''
        raise NotImplementedError('Abstract method')

    def get_state_names(self):
        '''
        Return the names of the recurrent states.
        It's required by apply_over_sequence method in order to allocate
        recurrent states for all steps with meaningful names.
        '''
        raise NotImplementedError('Abstract method')


class LSTMCell(RNNCell):

    def __init__(
        self,
        input_size,
        hidden_size,
        forget_bias,
        memory_optimization,
        name,
    ):
        super(LSTMCell, self).__init__(name)
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.forget_bias = float(forget_bias)
        self.memory_optimization = memory_optimization

    def _apply(
        self,
        model,
        input_t,
        seq_lengths,
        states,
        timestep,
    ):
        hidden_t_prev, cell_t_prev = states
        gates_t = model.FC(
            hidden_t_prev,
            self.scope('gates_t'),
            dim_in=self.hidden_size,
            dim_out=4 * self.hidden_size,
            axis=2,
        )
        model.net.Sum([gates_t, input_t], gates_t)
        hidden_t, cell_t = model.net.LSTMUnit(
            [
                hidden_t_prev,
                cell_t_prev,
                gates_t,
                seq_lengths,
                timestep,
            ],
            list(self.get_state_names()),
            forget_bias=self.forget_bias,
        )
        model.net.AddExternalOutputs(hidden_t, cell_t)
        if self.memory_optimization:
            self.recompute_blobs = [gates_t]
        return hidden_t, cell_t

    def get_input_params(self):
        return {
            'weights': self.scope('i2h') + '_w',
            'biases': self.scope('i2h') + '_b',
        }

    def get_recurrent_params(self):
        return {
            'weights': self.scope('gates_t') + '_w',
            'biases': self.scope('gates_t') + '_b',
        }

    def prepare_input(self, model, input_blob):
        return model.FC(
            input_blob,
            self.scope('i2h'),
            dim_in=self.input_size,
            dim_out=4 * self.hidden_size,
            axis=2,
        )

    def get_state_names(self):
        return (self.scope('hidden_t'), self.scope('cell_t'))

    def get_outputs_with_grads(self):
        return [0]

    def get_output_size(self):
        return self.hidden_size


def LSTM(model, input_blob, seq_lengths, initial_states, dim_in, dim_out,
         scope, outputs_with_grads=(0,), return_params=False,
         memory_optimization=False, forget_bias=0.0):
    '''
    Adds a standard LSTM recurrent network operator to a model.

    model: CNNModelHelper object new operators would be added to

    input_blob: the input sequence in a format T x N x D
    where T is sequence size, N - batch size and D - input dimention

    seq_lengths: blob containing sequence lengths which would be passed to
    LSTMUnit operator

    initial_states: a tupple of (hidden_input_blob, cell_input_blob)
    which are going to be inputs to the cell net on the first iteration

    dim_in: input dimention

    dim_out: output dimention

    outputs_with_grads : position indices of output blobs which will receive
    external error gradient during backpropagation

    return_params: if True, will return a dictionary of parameters of the LSTM

    memory_optimization: if enabled, the LSTM step is recomputed on backward step
                   so that we don't need to store forward activations for each
                   timestep. Saves memory with cost of computation.
    '''
    cell = LSTMCell(
        input_size=dim_in,
        hidden_size=dim_out,
        forget_bias=forget_bias,
        memory_optimization=memory_optimization,
        name=scope,
    )
    result = cell.apply_over_sequence(
        model=model,
        inputs=input_blob,
        seq_lengths=seq_lengths,
        initial_states=initial_states,
        outputs_with_grads=outputs_with_grads,
    )
    if return_params:
        result = list(result) + [{
            'input': cell.get_input_params(),
            'recurrent': cell.get_recurrent_params(),
        }]
    return tuple(result)


def GetLSTMParamNames():
    weight_params = ["input_gate_w", "forget_gate_w", "output_gate_w", "cell_w"]
    bias_params = ["input_gate_b", "forget_gate_b", "output_gate_b", "cell_b"]
    return {'weights': weight_params, 'biases': bias_params}


def InitFromLSTMParams(lstm_pblobs, param_values):
    '''
    Set the parameters of LSTM based on predefined values
    '''
    weight_params = GetLSTMParamNames()['weights']
    bias_params = GetLSTMParamNames()['biases']
    for input_type in param_values.keys():
        weight_values = [param_values[input_type][w].flatten() for w in weight_params]
        wmat = np.array([])
        for w in weight_values:
            wmat = np.append(wmat, w)
        bias_values = [param_values[input_type][b].flatten() for b in bias_params]
        bm = np.array([])
        for b in bias_values:
            bm = np.append(bm, b)

        weights_blob = lstm_pblobs[input_type]['weights']
        bias_blob = lstm_pblobs[input_type]['biases']
        cur_weight = workspace.FetchBlob(weights_blob)
        cur_biases = workspace.FetchBlob(bias_blob)

        workspace.FeedBlob(
            weights_blob,
            wmat.reshape(cur_weight.shape).astype(np.float32))
        workspace.FeedBlob(
            bias_blob,
            bm.reshape(cur_biases.shape).astype(np.float32))


def cudnn_LSTM(model, input_blob, initial_states, dim_in, dim_out,
               scope, recurrent_params=None, input_params=None,
               num_layers=1, return_params=False):
    '''
    CuDNN version of LSTM for GPUs.
    input_blob          Blob containing the input. Will need to be available
                        when param_init_net is run, because the sequence lengths
                        and batch sizes will be inferred from the size of this
                        blob.
    initial_states      tuple of (hidden_init, cell_init) blobs
    dim_in              input dimensions
    dim_out             output/hidden dimension
    scope               namescope to apply
    recurrent_params    dict of blobs containing values for recurrent
                        gate weights, biases (if None, use random init values)
                        See GetLSTMParamNames() for format.
    input_params        dict of blobs containing values for input
                        gate weights, biases (if None, use random init values)
                        See GetLSTMParamNames() for format.
    num_layers          number of LSTM layers
    return_params       if True, returns (param_extract_net, param_mapping)
                        where param_extract_net is a net that when run, will
                        populate the blobs specified in param_mapping with the
                        current gate weights and biases (input/recurrent).
                        Useful for assigning the values back to non-cuDNN
                        LSTM.
    '''
    with core.NameScope(scope):
        weight_params = GetLSTMParamNames()['weights']
        bias_params = GetLSTMParamNames()['biases']

        input_weight_size = dim_out * dim_in
        recurrent_weight_size = dim_out * dim_out
        input_bias_size = dim_out
        recurrent_bias_size = dim_out

        def init(layer, pname, input_type):
            if pname in weight_params:
                sz = input_weight_size if input_type == 'input' \
                    else recurrent_weight_size
            elif pname in bias_params:
                sz = input_bias_size if input_type == 'input' \
                    else recurrent_bias_size
            else:
                assert False, "unknown parameter type {}".format(pname)
            return model.param_init_net.UniformFill(
                [],
                "lstm_init_{}_{}_{}".format(input_type, pname, layer),
                shape=[sz])

        # Multiply by 4 since we have 4 gates per LSTM unit
        total_sz = 4 * num_layers * (
            input_weight_size + recurrent_weight_size + input_bias_size +
            recurrent_bias_size
        )

        weights = model.param_init_net.UniformFill(
            [], "lstm_weight", shape=[total_sz])

        model.params.append(weights)
        model.weights.append(weights)

        lstm_args = {
            'hidden_size': dim_out,
            'rnn_mode': 'lstm',
            'bidirectional': 0,  # TODO
            'dropout': 1.0,  # TODO
            'input_mode': 'linear',  # TODO
            'num_layers': num_layers,
            'engine': 'CUDNN'
        }

        param_extract_net = core.Net("lstm_param_extractor")
        param_extract_net.AddExternalInputs([input_blob, weights])
        param_extract_mapping = {}

        # Populate the weights-blob from blobs containing parameters for
        # the individual components of the LSTM, such as forget/input gate
        # weights and bises. Also, create a special param_extract_net that
        # can be used to grab those individual params from the black-box
        # weights blob. These results can be then fed to InitFromLSTMParams()
        for input_type in ['input', 'recurrent']:
            param_extract_mapping[input_type] = {}
            p = recurrent_params if input_type == 'recurrent' else input_params
            if p is None:
                p = {}
            for pname in weight_params + bias_params:
                for j in range(0, num_layers):
                    values = p[pname] if pname in p else init(j, pname, input_type)
                    model.param_init_net.RecurrentParamSet(
                        [input_blob, weights, values],
                        weights,
                        layer=j,
                        input_type=input_type,
                        param_type=pname,
                        **lstm_args
                    )
                    if pname not in param_extract_mapping[input_type]:
                        param_extract_mapping[input_type][pname] = {}
                    b = param_extract_net.RecurrentParamGet(
                        [input_blob, weights],
                        ["lstm_{}_{}_{}".format(input_type, pname, j)],
                        layer=j,
                        input_type=input_type,
                        param_type=pname,
                        **lstm_args
                    )
                    param_extract_mapping[input_type][pname][j] = b

        (hidden_input_blob, cell_input_blob) = initial_states
        output, hidden_output, cell_output, rnn_scratch, dropout_states = \
            model.net.Recurrent(
                [input_blob, cell_input_blob, cell_input_blob, weights],
                ["lstm_output", "lstm_hidden_output", "lstm_cell_output",
                 "lstm_rnn_scratch", "lstm_dropout_states"],
                seed=random.randint(0, 100000),  # TODO: dropout seed
                **lstm_args
            )
        model.net.AddExternalOutputs(
            hidden_output, cell_output, rnn_scratch, dropout_states)

    if return_params:
        param_extract = param_extract_net, param_extract_mapping
        return output, hidden_output, cell_output, param_extract
    else:
        return output, hidden_output, cell_output


class LSTMWithAttentionCell(RNNCell):

    def __init__(
        self,
        encoder_output_dim,
        encoder_outputs,
        decoder_input_dim,
        decoder_state_dim,
        name,
        attention_type,
        weighted_encoder_outputs,
        forget_bias,
        lstm_memory_optimization,
        attention_memory_optimization,
    ):
        super(LSTMWithAttentionCell, self).__init__(name)
        self.encoder_output_dim = encoder_output_dim
        self.encoder_outputs = encoder_outputs
        self.decoder_input_dim = decoder_input_dim
        self.decoder_state_dim = decoder_state_dim
        self.weighted_encoder_outputs = weighted_encoder_outputs
        self.encoder_outputs_transposed = None
        assert attention_type in [
            AttentionType.Regular,
            AttentionType.Recurrent,
        ]
        self.attention_type = attention_type
        self.lstm_memory_optimization = lstm_memory_optimization
        self.attention_memory_optimization = attention_memory_optimization

    def _apply(
        self,
        model,
        input_t,
        seq_lengths,
        states,
        timestep,
    ):
        (
            hidden_t_prev,
            cell_t_prev,
            attention_weighted_encoder_context_t_prev,
        ) = states

        gates_concatenated_input_t, _ = model.net.Concat(
            [hidden_t_prev, attention_weighted_encoder_context_t_prev],
            [
                self.scope('gates_concatenated_input_t'),
                self.scope('_gates_concatenated_input_t_concat_dims'),
            ],
            axis=2,
        )
        gates_t = model.FC(
            gates_concatenated_input_t,
            self.scope('gates_t'),
            dim_in=self.decoder_state_dim + self.encoder_output_dim,
            dim_out=4 * self.decoder_state_dim,
            axis=2,
        )
        model.net.Sum([gates_t, input_t], gates_t)

        hidden_t_intermediate, cell_t = model.net.LSTMUnit(
            [
                hidden_t_prev,
                cell_t_prev,
                gates_t,
                seq_lengths,
                timestep,
            ],
            ['hidden_t_intermediate', self.scope('cell_t')],
        )
        if self.attention_type == AttentionType.Recurrent:
            (
                attention_weighted_encoder_context_t,
                self.attention_weights_3d,
                attention_blobs,
            ) = apply_recurrent_attention(
                model=model,
                encoder_output_dim=self.encoder_output_dim,
                encoder_outputs_transposed=self.encoder_outputs_transposed,
                weighted_encoder_outputs=self.weighted_encoder_outputs,
                decoder_hidden_state_t=hidden_t_intermediate,
                decoder_hidden_state_dim=self.decoder_state_dim,
                scope=self.name,
                attention_weighted_encoder_context_t_prev=(
                    attention_weighted_encoder_context_t_prev
                ),
            )
        else:
            (
                attention_weighted_encoder_context_t,
                self.attention_weights_3d,
                attention_blobs,
            ) = apply_regular_attention(
                model=model,
                encoder_output_dim=self.encoder_output_dim,
                encoder_outputs_transposed=self.encoder_outputs_transposed,
                weighted_encoder_outputs=self.weighted_encoder_outputs,
                decoder_hidden_state_t=hidden_t_intermediate,
                decoder_hidden_state_dim=self.decoder_state_dim,
                scope=self.name,
            )
        hidden_t = model.Copy(hidden_t_intermediate, self.scope('hidden_t'))
        model.net.AddExternalOutputs(
            cell_t,
            hidden_t,
            attention_weighted_encoder_context_t,
        )
        if self.attention_memory_optimization:
            self.recompute_blobs.extend(attention_blobs)
        if self.lstm_memory_optimization:
            self.recompute_blobs.append(gates_t)

        return hidden_t, cell_t, attention_weighted_encoder_context_t

    def get_attention_weights(self):
        # [batch_size, encoder_length, 1]
        return self.attention_weights_3d

    def prepare_input(self, model, input_blob):
        if self.encoder_outputs_transposed is None:
            self.encoder_outputs_transposed = model.Transpose(
                self.encoder_outputs,
                self.scope('encoder_outputs_transposed'),
                axes=[1, 2, 0],
            )
        if self.weighted_encoder_outputs is None:
            self.weighted_encoder_outputs = model.FC(
                self.encoder_outputs,
                self.scope('weighted_encoder_outputs'),
                dim_in=self.encoder_output_dim,
                dim_out=self.encoder_output_dim,
                axis=2,
            )

        return model.FC(
            input_blob,
            self.scope('i2h'),
            dim_in=self.decoder_input_dim,
            dim_out=4 * self.decoder_state_dim,
            axis=2,
        )

    def get_state_names(self):
        return (
            self.scope('hidden_t'),
            self.scope('cell_t'),
            self.scope('attention_weighted_encoder_context_t'),
        )

    def get_outputs_with_grads(self):
        return [0, 4]

    def get_output_size(self):
        return self.decoder_state_dim + self.encoder_output_dim


def LSTMWithAttention(
    model,
    decoder_inputs,
    decoder_input_lengths,
    initial_decoder_hidden_state,
    initial_decoder_cell_state,
    initial_attention_weighted_encoder_context,
    encoder_output_dim,
    encoder_outputs,
    decoder_input_dim,
    decoder_state_dim,
    scope,
    attention_type=AttentionType.Regular,
    outputs_with_grads=(0, 4),
    weighted_encoder_outputs=None,
    lstm_memory_optimization=False,
    attention_memory_optimization=False,
    forget_bias=0.0,
):
    '''
    Adds a LSTM with attention mechanism to a model.

    The implementation is based on https://arxiv.org/abs/1409.0473, with
    a small difference in the order
    how we compute new attention context and new hidden state, similarly to
    https://arxiv.org/abs/1508.04025.

    The model uses encoder-decoder naming conventions,
    where the decoder is the sequence the op is iterating over,
    while computing the attention context over the encoder.

    model: CNNModelHelper object new operators would be added to

    decoder_inputs: the input sequence in a format T x N x D
    where T is sequence size, N - batch size and D - input dimention

    decoder_input_lengths: blob containing sequence lengths
    which would be passed to LSTMUnit operator

    initial_decoder_hidden_state: initial hidden state of LSTM

    initial_decoder_cell_state: initial cell state of LSTM

    initial_attention_weighted_encoder_context: initial attention context

    encoder_output_dim: dimension of encoder outputs

    encoder_outputs: the sequence, on which we compute the attention context
    at every iteration

    decoder_input_dim: input dimention (last dimension on decoder_inputs)

    decoder_state_dim: size of hidden states of LSTM

    attention_type: One of: AttentionType.Regular, AttentionType.Recurrent.
    Determines which type of attention mechanism to use.

    outputs_with_grads : position indices of output blobs which will receive
    external error gradient during backpropagation

    weighted_encoder_outputs: encoder outputs to be used to compute attention
    weights. In the basic case it's just linear transformation of
    encoder outputs (that the default, when weighted_encoder_outputs is None).
    However, it can be something more complicated - like a separate
    encoder network (for example, in case of convolutional encoder)

    lstm_memory_optimization: recompute LSTM activations on backward pass, so
                 we don't need to store their values in forward passes

    attention_memory_optimization: recompute attention for backward pass
    '''
    cell = LSTMWithAttentionCell(
        encoder_output_dim=encoder_output_dim,
        encoder_outputs=encoder_outputs,
        decoder_input_dim=decoder_input_dim,
        decoder_state_dim=decoder_state_dim,
        name=scope,
        attention_type=attention_type,
        weighted_encoder_outputs=weighted_encoder_outputs,
        forget_bias=forget_bias,
        lstm_memory_optimization=lstm_memory_optimization,
        attention_memory_optimization=attention_memory_optimization,
    )
    return cell.apply_over_sequence(
        model=model,
        inputs=decoder_inputs,
        seq_lengths=decoder_input_lengths,
        initial_states=(
            initial_decoder_hidden_state,
            initial_decoder_cell_state,
            initial_attention_weighted_encoder_context,
        ),
        outputs_with_grads=None,
    )


class MILSTMCell(LSTMCell):

    def _apply(
        self,
        model,
        input_t,
        seq_lengths,
        states,
        timestep,
    ):
        (
            hidden_t_prev,
            cell_t_prev,
        ) = states

        # hU^T
        # Shape: [1, batch_size, 4 * hidden_size]
        prev_t = model.FC(
            hidden_t_prev, self.scope('prev_t'), dim_in=self.hidden_size,
            dim_out=4 * self.hidden_size, axis=2)
        # defining MI parameters
        alpha = model.param_init_net.ConstantFill(
            [],
            [self.scope('alpha')],
            shape=[4 * self.hidden_size],
            value=1.0
        )
        beta1 = model.param_init_net.ConstantFill(
            [],
            [self.scope('beta1')],
            shape=[4 * self.hidden_size],
            value=1.0
        )
        beta2 = model.param_init_net.ConstantFill(
            [],
            [self.scope('beta2')],
            shape=[4 * self.hidden_size],
            value=1.0
        )
        b = model.param_init_net.ConstantFill(
            [],
            [self.scope('b')],
            shape=[4 * self.hidden_size],
            value=0.0
        )
        model.params.extend([alpha, beta1, beta2, b])
        # alpha * (xW^T * hU^T)
        # Shape: [1, batch_size, 4 * hidden_size]
        alpha_tdash = model.net.Mul(
            [prev_t, input_t],
            self.scope('alpha_tdash')
        )
        # Shape: [batch_size, 4 * hidden_size]
        alpha_tdash_rs, _ = model.net.Reshape(
            alpha_tdash,
            [self.scope('alpha_tdash_rs'), self.scope('alpha_tdash_old_shape')],
            shape=[-1, 4 * self.hidden_size],
        )
        alpha_t = model.net.Mul(
            [alpha_tdash_rs, alpha],
            self.scope('alpha_t'),
            broadcast=1,
            use_grad_hack=1
        )
        # beta1 * hU^T
        # Shape: [batch_size, 4 * hidden_size]
        prev_t_rs, _ = model.net.Reshape(
            prev_t,
            [self.scope('prev_t_rs'), self.scope('prev_t_old_shape')],
            shape=[-1, 4 * self.hidden_size],
        )
        beta1_t = model.net.Mul(
            [prev_t_rs, beta1],
            self.scope('beta1_t'),
            broadcast=1,
            use_grad_hack=1
        )
        # beta2 * xW^T
        # Shape: [batch_szie, 4 * hidden_size]
        input_t_rs, _ = model.net.Reshape(
            input_t,
            [self.scope('input_t_rs'), self.scope('input_t_old_shape')],
            shape=[-1, 4 * self.hidden_size],
        )
        beta2_t = model.net.Mul(
            [input_t_rs, beta2],
            self.scope('beta2_t'),
            broadcast=1,
            use_grad_hack=1
        )
        # Add 'em all up
        gates_tdash = model.net.Sum(
            [alpha_t, beta1_t, beta2_t],
            self.scope('gates_tdash')
        )
        gates_t = model.net.Add(
            [gates_tdash, b],
            self.scope('gates_t'),
            broadcast=1,
            use_grad_hack=1
        )
        # # Shape: [1, batch_size, 4 * hidden_size]
        gates_t_rs, _ = model.net.Reshape(
            gates_t,
            [self.scope('gates_t_rs'), self.scope('gates_t_old_shape')],
            shape=[1, -1, 4 * self.hidden_size],
        )

        hidden_t_intermediate, cell_t = model.net.LSTMUnit(
            [hidden_t_prev, cell_t_prev, gates_t_rs, seq_lengths, timestep],
            [self.scope('hidden_t_intermediate'), self.scope('cell_t')],
            forget_bias=self.forget_bias,
        )
        hidden_t = model.Copy(hidden_t_intermediate, self.scope('hidden_t'))
        model.net.AddExternalOutputs(
            cell_t,
            hidden_t,
        )
        if self.memory_optimization:
            self.recompute_blobs = [gates_t]
        return hidden_t, cell_t


def MILSTM(model, input_blob, seq_lengths, initial_states, dim_in, dim_out,
           scope, outputs_with_grads=(0,), memory_optimization=False,
           forget_bias=0.0):
    '''
    Adds MI flavor of standard LSTM recurrent network operator to a model.
    See https://arxiv.org/pdf/1606.06630.pdf

    model: CNNModelHelper object new operators would be added to

    input_blob: the input sequence in a format T x N x D
    where T is sequence size, N - batch size and D - input dimention

    seq_lengths: blob containing sequence lengths which would be passed to
    LSTMUnit operator

    initial_states: a tupple of (hidden_input_blob, cell_input_blob)
    which are going to be inputs to the cell net on the first iteration

    dim_in: input dimention

    dim_out: output dimention

    outputs_with_grads : position indices of output blobs which will receive
    external error gradient during backpropagation

    memory_optimization: if enabled, the LSTM step is recomputed on backward step
                   so that we don't need to store forward activations for each
                   timestep. Saves memory with cost of computation.
    '''
    cell = MILSTMCell(
        input_size=dim_in,
        hidden_size=dim_out,
        forget_bias=forget_bias,
        memory_optimization=memory_optimization,
        name=scope,
    )
    result = cell.apply_over_sequence(
        model=model,
        inputs=input_blob,
        seq_lengths=seq_lengths,
        initial_states=initial_states,
        outputs_with_grads=outputs_with_grads,
    )
    return tuple(result)


class MILSTMWithAttentionCell(LSTMWithAttentionCell):

    def _apply(
        self,
        model,
        input_t,
        seq_lengths,
        states,
        timestep,
    ):
        (
            hidden_t_prev,
            cell_t_prev,
            attention_weighted_encoder_context_t_prev,
        ) = states

        gates_concatenated_input_t, _ = model.net.Concat(
            [hidden_t_prev, attention_weighted_encoder_context_t_prev],
            [
                self.scope('gates_concatenated_input_t'),
                self.scope('_gates_concatenated_input_t_concat_dims'),
            ],
            axis=2,
        )
        # hU^T
        # Shape: [1, batch_size, 4 * hidden_size]
        prev_t = model.FC(
            gates_concatenated_input_t,
            self.scope('prev_t'),
            dim_in=self.decoder_state_dim + self.encoder_output_dim,
            dim_out=4 * self.decoder_state_dim,
            axis=2,
        )
        # defining MI parameters
        alpha = model.param_init_net.ConstantFill(
            [],
            [self.scope('alpha')],
            shape=[4 * self.decoder_state_dim],
            value=1.0
        )
        beta1 = model.param_init_net.ConstantFill(
            [],
            [self.scope('beta1')],
            shape=[4 * self.decoder_state_dim],
            value=1.0
        )
        beta2 = model.param_init_net.ConstantFill(
            [],
            [self.scope('beta2')],
            shape=[4 * self.decoder_state_dim],
            value=1.0
        )
        b = model.param_init_net.ConstantFill(
            [],
            [self.scope('b')],
            shape=[4 * self.decoder_state_dim],
            value=0.0
        )
        model.params.extend([alpha, beta1, beta2, b])
        # alpha * (xW^T * hU^T)
        # Shape: [1, batch_size, 4 * hidden_size]
        alpha_tdash = model.net.Mul(
            [prev_t, input_t],
            self.scope('alpha_tdash')
        )
        # Shape: [batch_size, 4 * hidden_size]
        alpha_tdash_rs, _ = model.net.Reshape(
            alpha_tdash,
            [self.scope('alpha_tdash_rs'), self.scope('alpha_tdash_old_shape')],
            shape=[-1, 4 * self.decoder_state_dim],
        )
        alpha_t = model.net.Mul(
            [alpha_tdash_rs, alpha],
            self.scope('alpha_t'),
            broadcast=1,
            use_grad_hack=1
        )
        # beta1 * hU^T
        # Shape: [batch_size, 4 * hidden_size]
        prev_t_rs, _ = model.net.Reshape(
            prev_t,
            [self.scope('prev_t_rs'), self.scope('prev_t_old_shape')],
            shape=[-1, 4 * self.decoder_state_dim],
        )
        beta1_t = model.net.Mul(
            [prev_t_rs, beta1],
            self.scope('beta1_t'),
            broadcast=1,
            use_grad_hack=1
        )
        # beta2 * xW^T
        # Shape: [batch_szie, 4 * hidden_size]
        input_t_rs, _ = model.net.Reshape(
            input_t,
            [self.scope('input_t_rs'), self.scope('input_t_old_shape')],
            shape=[-1, 4 * self.decoder_state_dim],
        )
        beta2_t = model.net.Mul(
            [input_t_rs, beta2],
            self.scope('beta2_t'),
            broadcast=1,
            use_grad_hack=1
        )
        # Add 'em all up
        gates_tdash = model.net.Sum(
            [alpha_t, beta1_t, beta2_t],
            self.scope('gates_tdash')
        )
        gates_t = model.net.Add(
            [gates_tdash, b],
            self.scope('gates_t'),
            broadcast=1,
            use_grad_hack=1
        )
        # # Shape: [1, batch_size, 4 * hidden_size]
        gates_t_rs, _ = model.net.Reshape(
            gates_t,
            [self.scope('gates_t_rs'), self.scope('gates_t_old_shape')],
            shape=[1, -1, 4 * self.decoder_state_dim],
        )

        hidden_t_intermediate, cell_t = model.net.LSTMUnit(
            [hidden_t_prev, cell_t_prev, gates_t_rs, seq_lengths, timestep],
            [self.scope('hidden_t_intermediate'), self.scope('cell_t')],
        )

        if self.attention_type == AttentionType.Recurrent:
            (
                attention_weighted_encoder_context_t,
                self.attention_weights_3d,
                self.recompute_blobs,
            ) = (
                apply_recurrent_attention(
                    model=model,
                    encoder_output_dim=self.encoder_output_dim,
                    encoder_outputs_transposed=self.encoder_outputs_transposed,
                    weighted_encoder_outputs=self.weighted_encoder_outputs,
                    decoder_hidden_state_t=hidden_t_intermediate,
                    decoder_hidden_state_dim=self.decoder_state_dim,
                    scope=self.name,
                    attention_weighted_encoder_context_t_prev=(
                        attention_weighted_encoder_context_t_prev
                    ),
                )
            )
        else:
            (
                attention_weighted_encoder_context_t,
                self.attention_weights_3d,
                self.recompute_blobs,
            ) = (
                apply_regular_attention(
                    model=model,
                    encoder_output_dim=self.encoder_output_dim,
                    encoder_outputs_transposed=self.encoder_outputs_transposed,
                    weighted_encoder_outputs=self.weighted_encoder_outputs,
                    decoder_hidden_state_t=hidden_t_intermediate,
                    decoder_hidden_state_dim=self.decoder_state_dim,
                    scope=self.name,
                )
            )
        hidden_t = model.Copy(hidden_t_intermediate, self.scope('hidden_t'))
        model.net.AddExternalOutputs(
            cell_t,
            hidden_t,
            attention_weighted_encoder_context_t,
        )
        return hidden_t, cell_t, attention_weighted_encoder_context_t